[BUG] RuntimeError: Invalid device string: 'bfloat16' with transformers v4.40.1 and save_strategy="epoch"

[OAHC2022]

While fine-tuning the unsloth/codellama-7b model using transformers v4.40.1 and setting save_strategy=epoch, I encountered the following error:

line 540, in LlamaModel_fast_forward
    inputs_embeds = inputs_embeds.to(self.config.torch_dtype)
RuntimeError: Invalid device string: 'bfloat16'

Upon examining the code, I identified the problematic line at this GitHub location:

# Embed positions
if inputs_embeds is None:
    inputs_embeds = self.embed_tokens(input_ids)

inputs_embeds = inputs_embeds.to(self.config.torch_dtype)

It appears that during the model's saving process in fine-tuning, the self.config.torch_dtype was incorrectly set as the string "bfloat16" instead of torch.bfloat16. Here's a simple fix I implemented:

# Embed positions
if inputs_embeds is None:
    inputs_embeds = self.embed_tokens(input_ids)

# My Modification
if self.config.torch_dtype == "bfloat16":
    self.config.torch_dtype = torch.bfloat16
inputs_embeds = inputs_embeds.to(self.config.torch_dtype)

[danielhanchen]

Oh thanks for that!! Will add your fix in! Thanks!

[OAHC2022]

Thank you!