Save LoRA weight in bfloat16

[thehs29]

My base model (llama3.2-3b-instruct) is in bfloat16, however, the lora adapter after training is float32, thus causing difference when merging adapter into a base model weight. So I would like to ask is there any way to have the LoRA adapter weights in bfloat16?

I train the model using FSDP + LoRA

[NanoCode012]

Hey, do you have your config available to share? Did you set bf16: true? What kind of system did you train on?

[thehs29]

about my training hardware, i used 8xA100

[NanoCode012]

How did you verify it was not in bf16?

Edit: Okay, I can repro this issue.

[thehs29]

i use safetensors.torch.load_file to load adapter_model.safetensors as state_dict, and print dtype of each tensor in that dict

[thehs29]

How did you verify it was not in bf16?

Edit: Okay, I can repro this issue.

so, any way to solve it?

P/S: actually, there is no big problem except that when merging weights the model becomes weird. I can still use lora adapter for inference, and, of course it'll be significantly slower

[NanoCode012]

I did an extra single GPU run to ensure it's not fsdp, but the adapter was also in fp32.

---

P/S: actually, there is no big problem except that when merging weights the model becomes weird.

If anything, this shouldn't be the issue? The model is not losing anymore precision had it have been a bf16 adapter when merged. Our merge cli axolotl merge-lora also casts it to torch_dtype which would've been bf16 in your case. Or did you use a custom merge script?

I can still use lora adapter for inference

I think you can tell peft to load in bf16/fp16 during inference? model = model.to(device, dtype=torch.bfloat16)

Perhaps, this is an upstream change that caused the adapters to now save in fp32.