Model weights saved in FP32 with FSDP when loaded model weights are in bf16
Author: shivam15s
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.45.1
- Platform: Linux-5.15.148.2-2.cm2-x86_64-with-glibc2.35
- Python version: 3.10.14
- Huggingface_hub version: 0.25.0
- Safetensors version: 0.4.5
- Accelerate version: 0.34.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.0+cu124 (True)
- Tensorflow version (GPU?): 2.16.2 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: yes
- Using GPU in script?: yes
- GPU type: NVIDIA A100-SXM4-80GB

### Who can help?

@ArthurZucker @muellerzr @SunMarc 


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

With FSDP enabled, weights are saved in float32, while they are saved in bfloat16 otherwise
```py
from datasets import load_dataset
from trl import SFTConfig, SFTTrainer
import transformers
import torch

dataset = load_dataset("openai/gsm8k", "main", split="train")

sft_config = SFTConfig(
    dataset_text_field="question",
    max_seq_length=4,
    per_device_train_batch_size=1,
    save_steps=4,
    output_dir="/tmp",
    bf16=True,
    fsdp="full_shard auto_wrap", # remove this line for seeing the other case
)
model = transformers.AutoModelForCausalLM.from_pretrained(
            "meta-llama/Meta-Llama-3-8B-Instruct",
            torch_dtype=torch.bfloat16,
        )

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    args=sft_config,
)
trainer.train()
```

### Expected behavior

Weights saved in bfloat16 with or without fsdp

Comments:
Hi, which trl version do you use?
Currently, I can't reproduce:

```python
from trl import SFTConfig, SFTTrainer
import json
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

dataset = load_dataset("openai/gsm8k", "main", split="train")

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)

training_args = SFTConfig(
    dataset_text_field="question",
    per_device_train_batch_size=1,
    save_steps=4,
    output_dir="tmp",
    max_steps=4,
    fsdp="full_shard auto_wrap",
)

model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)

trainer = SFTTrainer(model=model, train_dataset=dataset, args=training_args)
trainer.train()

with open("tmp/checkpoint-4/config.json", "r") as f:
    config = json.load(f)

print(config["torch_dtype"])  # bfloat16
```

I'm using `trl Version: 0.11.0`
I think you need to set `bf16=True` in your `SFTConfig` to reproduce

It seems this behavior is expected based on this reference: https://github.com/huggingface/accelerate/blob/5060574827ffcf1055179896cbb4f150caa3aec5/src/accelerate/accelerator.py#L1521. 
I’d appreciate your thoughts on this. Additionally, do you know how we can save directly in bf16 in this context?

cc @fabianlim @muellerzr 

@shivam15s if you want to save in bf16, then you should not be training in mixed precision , so set `bf16 == False`. 
- the purpose of that change i made is for folks who have weights in bf16, but want to train in mixed precision, so therefore, the FSDP flat params will be upcasted. then in this case, the checkpoint weights surely must be in fp32, otherwise resumption will be impossible
- so if you want to save checkpoints and weights in bf16, then you surely cannot use mixed precision, for the above reason that you will never be able to resume your training properly, so then you must turn off mixed precision in this use case

This behavior is a generalization that of DeepSpeed, whereby in DS mixed-precision is default,  and the weights will always be saved in fp32. 

Gotcha, thanks for the detailed reply! I appreciate the help. Closing the issue.

@qgallouedec @fabianlim 
Reopening this issue because I am getting errors loading the saved model using `AutoModelForCausalLM.from_pretrained`
Here's the error I get: 
```
RuntimeError: Error(s) in loading state_dict for LlamaForCausalLM:
        size mismatch for model.embed_tokens.weight: copying a param with shape torch.Size([131334656]) from checkpoint, the shape in current model is torch.Size([128256, 4096]).
        size mismatch for model.norm.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([4096]).
        size mismatch for lm_head.weight: copying a param with shape torch.Size([0]) from checkpoint, the shape in current model is torch.Size([128256, 4096]).
        You may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method.
```
Here's the minimal reproducible code for it: 
```py
from trl import SFTConfig, SFTTrainer
import json
from datasets import load_dataset
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

dataset = load_dataset("openai/gsm8k", "main", split="train")

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)

training_args = SFTConfig(
    dataset_text_field="question",
    per_device_train_batch_size=1,
    bf16=True,
    save_strategy="no",
    output_dir="./tmp",
    max_steps=4,
    fsdp="full_shard auto_wrap",
)

model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)

trainer = SFTTrainer(model=model, train_dataset=dataset, args=training_args)
trainer.train()

model_save_path = os.path.join(training_args.output_dir, "model")
print(model_save_path)
if trainer.is_fsdp_enabled:
    trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")
trainer.save_model(model_save_path)
```
Would like to hear your thoughts on how to resolve this

@shivam15s the problem you pointed out is a different one and is related to this PR introduced in 0.34 https://github.com/huggingface/accelerate/pull/2959

It has already been reverted in main, but there has not yet been a release for it. In the interim  you can try to downgrade `pip install "accelerate<0.34"`

It'll be released this week once we fully launch out 1.0 :)

Confirmed that this was the reason. Thanks @fabianlim 👍