Init on meta device and then materialize on gpu leads to very large errors
Author: fingertap
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.47.1
- Platform: Linux-5.4.0-153-generic-x86_64-with-glibc2.35
- Python version: 3.10.16
- Huggingface_hub version: 0.27.1
- Safetensors version: 0.5.1
- Accelerate version: 1.2.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.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 distributed or parallel set-up in script?: No
- Using GPU in script?: Yes
- GPU type: NVIDIA A800-SXM4-80GB

### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

### Issue

I prepared a [repo](https://github.com/fingertap/meta_init_issue) to reproduce this issue. 

The key is that the language loss produced by the model using the following initialization pipeline is very large compared to directly initialize it on the GPU with `from_pretrained`. The pipeline to create large-error model

1. init the model with `init_empty_weights`
2. materialize the model with `to_empty`
3. load in the weights with `load_state_dict` from a model that is created by `from_pretrained`.

### Motivation

The reason why I want to initialize the model in this way is that I need to support load large models into GPU using FSDP. I found that by passing a meta-device model to FSDP with `param_init_fn` being `to_empty`, the language loss is large. If I pass a cuda model, it is ok. Later I found that this is not the problem of FSDP. Using `to_empty` on single-GPU case will still lead to large loss. In the era of large language models, this issue is **VERY CRITICAL** for scaling up to larger models.


### Expected behavior

A normal loss after loading the state dict, no matter how the model is initialized (should work after `to_empty`). In my repro example, the loss shoule be ~0.15. At least it should < 1.

Comments:
@SunMarc Can you comment on this?

Thanks for creating a repo for us to reproduce. A few comments from the code:
- I see that you are not using flash attention 2 for model_meta_accelerate
- if you do to_empty or with torch.device("meta"), you will lose the non persistant buffer that are not saved in the state dict, thus, you will end up with an incomplete model. You should use init_empty_weights instead from accelerate 
Can you compare the model params and buffers of each model to see where are the differences ? That will help a lot I think


Yes! I just read the `from_pretrained` and debug into it. The cause is exactly the buffer getting cleared. This works as expected.

```python
    with init_empty_weights():
            from transformers import Qwen2ForCausalLM, Qwen2Config
            config = Qwen2Config.from_pretrained(self.model_path, torch_dtype=self.dtype)
            model = Qwen2ForCausalLM(config)
            model.to(torch.bfloat16)
            
        another_model = AutoModelForCausalLM.from_pretrained(
            self.model_path,
            torch_dtype=self.dtype,
        )
        state_dict = another_model.state_dict()
        from transformers.modeling_utils import _load_state_dict_into_meta_model
        new_error_msgs, offload_index, state_dict_index = _load_state_dict_into_meta_model(
            model,
            state_dict,
            "",
            expected_keys=list(state_dict.keys()),
            dtype=torch.bfloat16
        )
        
        return model
```
Just before closing this issue, is there some tools that I have missed for init modules in FSDP so that I can avoid the `to_empty` in FSDP?

One solution that occurs to me is to first copy the buffers out, call `to_empty`, and then copy it back.

that could be one solution indeed. otherwise, you can also initialize them again by creating a `init_buffers` method. 

Thanks a ton!