Repository: huggingface/transformers
Pull request: fix FSDP loading with meta devices

# What does this PR do?

supersedes #44446

on `main`, when loading to cpu and using meta devices for non-rank0 processes, it now re-initializes weights on those processes as well as uses more CPU memory. In testing with loading llama3-8b.

main; both on CPU, uses 16GB system RAM, slow to load, re-inits weights on rank1
#44446: rank0 on CPU, rank1 on meta, uses 1.5GB system RAM
v4.57.6, both on CPU, uses 1.5GB system RAM
this PR, both on CPU, uses 1.5GB system RAM, same behavior and training loss as main and v4.57.6

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker @Cyrilvallez
- vision models: @yonigozlan @molbap
- audio models: @eustlb @ebezzam @vasqu
- multimodal models: @zucchini-nlp
- graph models: @clefourrier

Library:

- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- continuous batching: @remi-or @ArthurZucker @McPatate
- pipelines: @Rocketknight1
- tokenizers: @ArthurZucker and @itazap
- trainer: @SunMarc
- attention: @vasqu @ArthurZucker @CyrilVallez
- model loading (from pretrained, etc): @CyrilVallez
- distributed: @3outeille @ArthurZucker
- CIs: @ydshieh

Integrations:

- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization: @SunMarc
- kernels: @drbh
- peft: @BenjaminBossan @githubnemo

Devices/Backends:

- AMD ROCm: @ivarflakstad
- Intel XPU: @IlyasMoutawwakil
- Ascend NPU: @ivarflakstad 

Documentation: @stevhliu

Research projects are not maintained and should be taken as is.

 -->


Changed files:
filename: src/transformers/modeling_utils.py
patch: @@ -4548,8 +4548,23 @@ def _initialize_missing_keys(self, is_quantized: bool) -> None:
         `_initialize_weights`. Indeed, since the corresponding weights are missing from the state dict, they will not be replaced and need to
         be initialized correctly (i.e. weight initialization distribution).
 
-        Params that are not missing have the `is_hf_initialized` flag.
+        Also marks non-missing params/buffers with `_is_hf_initialized` and propagates this flag to modules,
+        so that `_initialize_weights` can skip fully-initialized modules entirely.
         """
+        if is_fsdp_enabled() and not is_local_dist_rank_0():
+            # Handle FSDP edge case when using cpu ram efficient loading to ensure it is marked as initialized
+            # since it will get its weights broadcasted from rank0
+            # We actually need to do that only because we want to re-initialize non-persistent buffers with correct values.
+            # Everything else in the state_dict will be gathered from rank0, so we don't need re-initialization.
+            # We could simply early return after buffer inits if we had a way to init only the non-persistent buffers
+            for key in self.state_dict():
+                try:
+                    param_or_buffer = self.get_parameter_or_buffer(key)
+                    param_or_buffer._is_hf_initialized = True
+                except AttributeError:
+                    pass  # may happen when handling pre-quantized weights
+            self._is_hf_initialized = True
+
         # This will only initialize submodules that are not marked as initialized by the line above.
         if is_deepspeed_zero3_enabled() and not is_quantized:
             import deepspeed

filename: tests/utils/test_modeling_utils.py
patch: @@ -2361,6 +2361,127 @@ def create_bidirectional_mask_with_kwargs(
         torch.testing.assert_close(reference, with_config_only)
 
 
+@require_torch
+class InitializeMissingKeysTest(unittest.TestCase):
+    """Tests for _initialize_missing_keys to prevent regressions in FSDP non-rank-0 weight initialization.
+
+    On FSDP non-rank-0 processes, weights are loaded on meta device and then moved to CPU via
+    _move_missing_keys_from_meta_to_device. The _initialize_missing_keys method must mark all
+    params as _is_hf_initialized so that guarded init functions (init.normal_, init.zeros_, etc.)
+    skip them, preventing expensive and unnecessary re-initialization.
+    """
+
+    def _clear_init_flags(self, model):
+        """Clear all _is_hf_initialized flags from params, buffers, and modules."""
+        for module in model.modules():
+            if hasattr(module, "_is_hf_initialized"):
+                delattr(module, "_is_hf_initialized")
+        for param in model.parameters():
+            if hasattr(param, "_is_hf_initialized"):
+                delattr(param, "_is_hf_initialized")
+        for buffer in model.buffers():
+            if hasattr(buffer, "_is_hf_initialized"):
+                delattr(buffer, "_is_hf_initialized")
+
+    @staticmethod
+    def _fsdp_non_rank0_patches():
+        """Return a context manager that mocks FSDP as enabled on a non-rank-0 process."""
+        from contextlib import ExitStack
+
+        stack = ExitStack()
+        stack.enter_context(patch("transformers.modeling_utils.is_fsdp_enabled", return_value=True))
+        stack.enter_context(patch("transformers.modeling_utils.is_local_dist_rank_0", return_value=False))
+        return stack
+
+    def test_initialize_missing_keys_marks_all_params(self):
+        """On FSDP non-rank-0, all params/buffers should get _is_hf_initialized=True."""
+        model = BaseModel(PreTrainedConfig())
+        self._clear_init_flags(model)
+
+        with self._fsdp_non_rank0_patches():
+            model._initialize_missing_keys(is_quantized=False)
+
+        for name, param in model.named_parameters():
+            self.assertTrue(
+                getattr(param, "_is_hf_initialized", False),
+                f"param {name} should be marked as initialized",
+            )
+
+    def test_initialize_missing_keys_sets_model_flag(self):
+        """On FSDP non-rank-0, the model itself should get _is_hf_initialized=True."""
+        model = BaseModel(PreTrainedConfig())
+        self._clear_init_flags(model)
+
+        with self._fsdp_non_rank0_patches():
+            model._initialize_missing_keys(is_quantized=False)
+
+        self.assertTrue(getattr(model, "_is_hf_initialized", False))
+
+    def test_initialize_missing_keys_no_reinit_of_marked_params(self):
+        """Guarded init functions should not modify params that have _is_hf_initialized=True."""
+        model = BaseModel(PreTrainedConfig())
+        self._clear_init_flags(model)
+
+        # Snapshot values before _initialize_missing_keys
+        pre_values = {name: param.clone() for name, param in model.named_parameters()}
+
+        with self._fsdp_non_rank0_patches():
+            model._initialize_missing_keys(is_quantized=False)
+
+        # Params should not have changed (guarded inits skip marked params)
+        for name, param in model.named_parameters():
+            torch.testing.assert_close(param, pre_values[name], msg=f"param {name} should not be re-initialized")
+
+    def test_move_missing_keys_fsdp_non_rank0_moves_meta_to_cpu(self):
+        """FSDP non-rank-0 path should move all params/buffers from meta to CPU."""
+        with torch.device("meta"):
+            model = BaseModel(PreTrainedConfig())
+
+        # Verify everything starts on meta
+        for param in model.parameters():
+            self.assertEqual(param.device, torch.device("meta"))
+
+        with (
+            patch("transformers.modeling_utils.is_fsdp_enabled", return_value=True),
+            patch("transformers.modeling_utils.is_local_dist_rank_0", return_value=False),
+        ):
+            model._move_missing_keys_from_meta_to_device(
+                missing_keys=set(), device_map=None, device_mesh=None, hf_quantizer=None
+            )
+
+        # All params should now be on CPU
+        for name, param in model.named_parameters():
+            self.assertEqual(param.device, torch.device("cpu"), f"param {name} should be on CPU after FSDP move")
+
+    def test_fsdp_non_rank0_end_to_end_no_reinit(self):
+        """End-to-end: FSDP non-rank-0 move + _initialize_missing_keys should not re-init params."""
+        with torch.device("meta"):
+            model = BaseModel(PreTrainedConfig())
+
+        with self._fsdp_non_rank0_patches():
+            model._move_missing_keys_from_meta_to_device(
+                missing_keys=set(), device_map=None, device_mesh=None, hf_quantizer=None
+            )
+
+            # After move: params are empty CPU tensors. Snapshot values.
+            pre_init_values = {name: param.clone() for name, param in model.named_parameters()}
+
+            model._initialize_missing_keys(is_quantized=False)
+
+        # All params should be marked as initialized
+        for name, param in model.named_parameters():
+            self.assertTrue(
+                getattr(param, "_is_hf_initialized", False), f"param {name} should be marked as initialized"
+            )
+
+        # Model should be marked as initialized
+        self.assertTrue(getattr(model, "_is_hf_initialized", False), "model should be marked as initialized")
+
+        # Param values should NOT have changed (no re-initialization)
+        for name, param in model.named_parameters():
+            torch.testing.assert_close(param, pre_init_values[name], msg=f"param {name} should not be re-initialized")
+
+
 @slow
 @require_torch
 class ModelOnTheFlyConversionTester(unittest.TestCase):