Init on meta device and then materialize on gpu leads to very large errors

[fingertap]

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
• 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

Issue

I prepared a repo 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.

[fingertap]

@SunMarc Can you comment on this?

[SunMarc]

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

[fingertap]

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

    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?

[fingertap]

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

[SunMarc]

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

[fingertap]

Thanks a ton!