fix(checkpoint): load base weights under DDP and MegatronFSDP

nemo_automodel/_transformers/infrastructure.py

@@ -593,21 +593,23 @@ def apply_model_infrastructure(
 
     checkpoint_already_loaded = False
     if load_before_shard:
-        if is_meta_device:
-            lora_a_init = getattr(peft_config, "lora_A_init", None)
-            checkpointer.initialize_model_weights(model, device, peft_init_method=lora_a_init)
+        if weights_already_loaded:
+            # HF's from_pretrained already populated the weights during model init.
+            # Still call load_base_model with load_base_model=False to
+            # handle weight tying
+            checkpointer.load_base_model(model, device, cache_dir, pretrained_model_name_or_path, load_base_model=False)
+        else:
+            # Only meta-device models need their parameter shells materialized first.
+            if is_meta_device:
+                lora_a_init = getattr(peft_config, "lora_A_init", None)
+                checkpointer.initialize_model_weights(model, device, peft_init_method=lora_a_init)
             checkpointer.load_base_model(
                 model,
                 device,
                 cache_dir,
                 pretrained_model_name_or_path,
                 load_base_model=load_base_model,
             )
-        else:
-            # Non-meta models already have weights from from_pretrained.
-            # Still call load_base_model with load_base_model=False to
-            # handle weight tying
-            checkpointer.load_base_model(model, device, cache_dir, pretrained_model_name_or_path, load_base_model=False)
         checkpoint_already_loaded = True
 
     # hold a list copy of the model state dict keys before any parallelization. To be used during checkpoint saving in safetensors format.

tests/unit_tests/_transformers/test_auto_model.py

@@ -1727,6 +1727,43 @@ def test_meta_tensor_not_implemented_error_retries_without_meta_device_on_hf_pat
             assert result is sentinel_model
             assert mock_init.call_count == 2
 
+    def test_custom_model_under_ddp_still_needs_its_checkpoint(self):
+        """A MODEL_ARCH_MAPPING model under DDP reaches infrastructure unloaded and off meta.
+
+        ``DDPManager`` is excluded from meta-device init, and custom model constructors
+        only build the architecture, so ``apply_model_infrastructure`` has to be told the
+        weights are still missing. If either flag is wrong the model enters training
+        randomly initialized and nothing is raised.
+        """
+        build_kwargs, mock_config = self._make_build_kwargs()
+        build_kwargs["is_hf_model"] = False
+        dummy_manager_cls = type("DummyManager", (), {})
+        build_kwargs["model_wrapper"] = dummy_manager_cls()
+        sentinel_model = MagicMock()
+        captured = {}
+
+        def capture(**kwargs):
+            captured.update(kwargs)
+            return sentinel_model
+
+        with (
+            patch("nemo_automodel._transformers.auto_model.DDPManager", dummy_manager_cls),
+            patch("nemo_automodel._transformers.auto_model._init_model", return_value=(True, sentinel_model)),
+            patch("nemo_automodel._transformers.auto_model.get_world_size_safe", return_value=1),
+            patch(
+                "nemo_automodel._transformers.capabilities.attach_capabilities_and_validate",
+                return_value=sentinel_model,
+            ),
+            patch("nemo_automodel._transformers.auto_model.apply_model_infrastructure", side_effect=capture),
+            patch("nemo_automodel._transformers.auto_model.get_hf_config", return_value=mock_config),
+            patch("nemo_automodel._transformers.auto_model._maybe_dequantize_fp8_for_peft", return_value=False),
+            patch("torch.cuda.current_device", return_value=0),
+        ):
+            _BaseNeMoAutoModelClass._build_model(mock_config, **build_kwargs)
+
+        assert captured["is_meta_device"] is False
+        assert captured["weights_already_loaded"] is False
+
 
 class TestNeMoAutoModelForMultimodalLM:
     """Tests for the NeMoAutoModelForMultimodalLM class and its exports."""

tests/unit_tests/_transformers/test_infrastructure.py

@@ -569,7 +569,9 @@ def test_skips_rotary_fix_when_not_needed(self):
 # =============================================================================
 
 
-def _run_apply_model_infrastructure_load_before_shard(*, peft_config=None):
+def _run_apply_model_infrastructure_load_before_shard(
+    *, peft_config=None, is_meta_device=True, weights_already_loaded=False
+):
     """Helper that invokes apply_model_infrastructure with load_before_shard=True."""
     from nemo_automodel._transformers.infrastructure import apply_model_infrastructure
 
@@ -588,12 +590,13 @@ def _run_apply_model_infrastructure_load_before_shard(*, peft_config=None):
 
         result = apply_model_infrastructure(
             model=model,
-            is_meta_device=True,
+            is_meta_device=is_meta_device,
             device=torch.device("cpu"),
             load_base_model=True,
             peft_config=peft_config,
             pretrained_model_name_or_path="test/model",
             cache_dir="/tmp/cache",
+            weights_already_loaded=weights_already_loaded,
         )
 
         return result, mock_ckpt, model
@@ -632,6 +635,76 @@ def test_load_before_shard_does_not_call_load_base_model_with_peft_init_method(s
         _, kwargs = mock_ckpt.load_base_model.call_args
         assert "peft_init_method" not in kwargs
 
+    def test_load_before_shard_loads_checkpoint_when_init_left_weights_unloaded(self):
+        """A non-meta model whose init did not load weights must still read the checkpoint.
+
+        ``is_meta_device`` is False for every model built under DDP or MegatronFSDP,
+        including AutoModel's own implementations, whose constructor only creates the
+        architecture. Gating the read on it left those models randomly initialized.
+        """
+        _, mock_ckpt, model = _run_apply_model_infrastructure_load_before_shard(
+            is_meta_device=False, weights_already_loaded=False
+        )
+
+        mock_ckpt.load_base_model.assert_called_once_with(
+            model, torch.device("cpu"), "/tmp/cache", "test/model", load_base_model=True
+        )
+        # Nothing is on meta, so there are no parameter shells to materialize.
+        mock_ckpt.initialize_model_weights.assert_not_called()
+
+    def test_load_before_shard_skips_checkpoint_when_init_already_loaded_weights(self):
+        """HF's from_pretrained already populated the weights; only re-tie, do not re-read."""
+        _, mock_ckpt, model = _run_apply_model_infrastructure_load_before_shard(
+            is_meta_device=False, weights_already_loaded=True
+        )
+
+        mock_ckpt.load_base_model.assert_called_once_with(
+            model, torch.device("cpu"), "/tmp/cache", "test/model", load_base_model=False
+        )
+        mock_ckpt.initialize_model_weights.assert_not_called()
+
+
+def test_load_before_shard_populates_unloaded_model_from_checkpoint(tmp_path):
+    """End-to-end guard on CPU: the model must end up holding the checkpoint's values.
+
+    The mocked tests above pin the call; this one pins the outcome, using the real
+    Checkpointer and a real safetensors checkpoint whose every tensor is 0.5.
+    """
+    from transformers import LlamaConfig
+    from transformers import LlamaForCausalLM as HFLlamaForCausalLM
+
+    from nemo_automodel._transformers.infrastructure import apply_model_infrastructure
+    from nemo_automodel.components.models.llama.model import LlamaForCausalLM
+
+    config = LlamaConfig(
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        vocab_size=256,
+        max_position_embeddings=128,
+        tie_word_embeddings=False,
+    )
+    reference = HFLlamaForCausalLM(config)
+    with torch.no_grad():
+        for param in reference.parameters():
+            param.fill_(0.5)
+    reference.save_pretrained(tmp_path, safe_serialization=True)
+
+    config.torch_dtype = torch.float32
+    model = apply_model_infrastructure(
+        model=LlamaForCausalLM(config),
+        is_meta_device=False,
+        device=torch.device("cpu"),
+        load_base_model=True,
+        pretrained_model_name_or_path=str(tmp_path),
+        weights_already_loaded=False,
+    )
+
+    for name, param in model.named_parameters():
+        assert torch.equal(param.detach(), torch.full_like(param, 0.5)), f"{name} was not loaded"
+
 
 # =============================================================================
 # Tests for from_config load_base_model kwarg forwarding