Honor per-param _is_hf_initialized flags for built-in models in _initialize_weights

src/transformers/modeling_utils.py

@@ -4771,6 +4771,24 @@ def _move_missing_keys_from_meta_to_device(
             value = torch.empty_like(buffer, device=buffer_device)
             _load_parameter_into_model(self, key, value)
 
+    def _mark_fully_loaded_submodules_as_initialized(self) -> None:
+        """Set the module-level `_is_hf_initialized` flag on every submodule (including `self`) whose parameters and
+        buffers are all already flagged as initialized, so `_initialize_weights` skips `_init_weights` for them.
+
+        Used on non-rank-0 FSDP/ZeRO processes, where the loaded params/buffers are flagged (they are broadcast from
+        rank 0) but their enclosing modules are not, which would otherwise re-run `_init_weights` on every submodule
+        (wasteful in general, and very costly on accelerators where `normal_` is far slower than on CUDA). Modules with
+        unflagged buffers (e.g. non-persistent buffers absent from the state dict) are left untouched so they are still
+        re-initialized with correct values.
+        """
+        for submodule in self.modules():
+            params_ready = all(getattr(p, "_is_hf_initialized", False) for p in submodule.parameters(recurse=False))
+            buffers_ready = all(
+                getattr(b, "_is_hf_initialized", False) for b in submodule.buffers(recurse=False) if b is not None
+            )
+            if params_ready and buffers_ready:
+                submodule._is_hf_initialized = True
+
     def _initialize_missing_keys(self, is_quantized: bool) -> None:
         """
         Initialize the missing keys (keys that are part of the model parameters, but were NOT found in the loaded state dicts), according to
@@ -4792,7 +4810,7 @@ def _initialize_missing_keys(self, is_quantized: bool) -> None:
                     param_or_buffer._is_hf_initialized = True
                 except AttributeError:
                     pass  # may happen when handling pre-quantized weights
-            self._is_hf_initialized = True
+            self._mark_fully_loaded_submodules_as_initialized()
 
         # This will only initialize submodules that are not marked as initialized by the line above.
         if is_deepspeed_zero3_enabled() and not is_quantized:

tests/utils/test_modeling_utils.py

@@ -475,6 +475,37 @@ def test_model_from_pretrained(self):
         self.assertEqual(model.config.output_hidden_states, True)
         self.assertEqual(model.config, config)
 
+    def test_mark_fully_loaded_submodules_as_initialized(self):
+        # Regression test for #47427: on non-rank-0 FSDP/ZeRO processes the loaded params/buffers are flagged but
+        # their enclosing modules are not, so `_init_weights` would redundantly re-run on every submodule. A module
+        # whose params/buffers are all flagged must get the module-level flag (so `_initialize_weights` skips it),
+        # while a module with an unflagged buffer (e.g. a non-persistent buffer) must be left to re-initialize.
+        config = BertConfig(
+            hidden_size=32, num_hidden_layers=1, num_attention_heads=2, intermediate_size=37, vocab_size=99
+        )
+        model = BertModel(config)
+        for submodule in model.modules():
+            submodule._is_hf_initialized = False
+
+        flagged = model.encoder.layer[0].attention.self.query
+        for param in flagged.parameters(recurse=False):
+            param._is_hf_initialized = True
+
+        with_unflagged_buffer = model.embeddings
+        for param in with_unflagged_buffer.parameters(recurse=False):
+            param._is_hf_initialized = True
+        # position_ids is a (non-persistent) buffer left unflagged.
+
+        model._mark_fully_loaded_submodules_as_initialized()
+
+        self.assertTrue(flagged._is_hf_initialized)
+        self.assertFalse(with_unflagged_buffer._is_hf_initialized)
+
+        # And a module marked as initialized is then skipped by `_initialize_weights`.
+        with mock.patch.object(type(model), "_init_weights") as mocked_init:
+            model._initialize_weights(flagged)
+        mocked_init.assert_not_called()
+
     def test_model_from_pretrained_subfolder(self):
         config = BertConfig.from_pretrained("hf-internal-testing/tiny-random-bert")
         model = BertModel(config)
@@ -1747,7 +1778,9 @@ def test_unexpected_keys_warnings(self):
                 with CaptureLogger(logger) as cl:
                     _, loading_info = ModelWithHead.from_pretrained(tmp_dir, output_loading_info=True)
             # Will be colored if terminal is interactive
-            expected_output = "added_key | UNEXPECTED" if sys.stdout.isatty() else "added_key | UNEXPECTED"
+            expected_output = (
+                "added_key | \x1b[38;5;208mUNEXPECTED" if sys.stdout.isatty() else "added_key | UNEXPECTED"
+            )
             self.assertIn(expected_output, cl.out)
             self.assertEqual(loading_info["unexpected_keys"], {"added_key"})