I got a float32 model when I train the model using FSDP and bf16. And There is precision error when I manually transfer the output model to bf16 model. 
Author: jiangjiadi
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.41.2
- Platform: Linux-4.9.151-015.ali3000.alios7.x86_64-x86_64-with-glibc2.17
- Python version: 3.8.18
- Huggingface_hub version: 0.23.2
- Safetensors version: 0.4.1
- Accelerate version: 0.30.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

 @muellerzr @SunMarc

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Set the following args as you train the model.
--bf16 True \
--fsdp "full_shard auto_wrap" \
--fsdp_config "./fsdp_config.json" 


### Expected behavior

We expect the output model to be in bfloat16 rather than float32. Upon investigation, I found that during the accelerator.prepare_model phase, Accelerate upcasts the parameters of the FSDP model when using mixed precision, causing the model parameter weights to become float32. See the traning logs:

> /opt/conda/lib/python3.8/site-packages/accelerate/accelerator.py:1558: UserWarning: Upcasted low precision parameters in Qwen2ForCausalLM because mixed precision turned on in FSDP. Affects: model.embed_tokens.weight, model.norm.weight, lm_head.weight.
/opt/conda/lib/python3.8/site-packages/accelerate/accelerator.py:1558: UserWarning: Upcasted low precision parameters in Qwen2DecoderLayer because mixed precision turned on in FSDP. Affects: self_attn.q_proj.weight, self_attn.q_proj.bias, self_attn.k_proj.weight, self_attn.k_proj.bias, self_attn.v_proj.weight, self_attn.v_proj.bias, self_attn.o_proj.weight, mlp.gate_proj.weight, mlp.up_proj.weight, mlp.down_proj.weight, input_layernorm.weight, post_attention_layernorm.weight.
/opt/conda/lib/python3.8/site-packages/accelerate/accelerator.py:1564: UserWarning: FSDP upcast of low precision parameters may affect the precision of model checkpoints.


This is somewhat confusing to me, as I don't understand why mixed precision training is enabled when training with bfloat16. 

Moreover, when I observed the parameters of the resulting float32 model, I found that directly converting the model parameters to bfloat16 introduces precision errors. Specifically, model_tensor.to(torch.bfloat16).to(torch.float32) - model_tensor != 0.0, where model_tensor is one of the float32 tensors in the model.

Comments:
Get this bug as well!

It's not a bug, we keep the base weights in fp32 for better end-model performance. Read this doc for more: (and it's why that error message is a thing) https://huggingface.co/docs/accelerate/concept_guides/fsdp_and_deepspeed#on-differences-in-data-precision-handling

I see, thanks for the answer!

Thanks!

@jiangjiadi Hello, have you taken any methods to solve the problem which can not save the trained checkpoints according to the user demand of precision? I also found that I always got a fp32 model using fsdp training, while I want to get a bf16 model if I set bf16 True.

@canghongjian You can overwrite the `save_model` function in Trainer to implement your own save logic.

I'm getting same warnings.