[DSV3] Fix the ckpt loading issue when no MoE layer on the mtp rank (#3315)

Signed-off-by: Gao Deng <gdeng@login-lyris02.lyris.clusters.nvidia.com>
Signed-off-by: Gao <gdeng@nvidia.com>
Co-authored-by: Gao Deng <gdeng@login-lyris02.lyris.clusters.nvidia.com>
Co-authored-by: Huy Vu <86480512+huvunvidia@users.noreply.github.com>

Author: 3 people

src/megatron/bridge/models/gpt_provider.py

@@ -361,7 +361,17 @@ def mtp_block_spec(config: "GPTModelProvider", vp_stage: Optional[int] = None) -
         if hasattr(spec, "layer_specs") and len(spec.layer_specs) == 0:
             # Get the decoder layer spec explicitly if no decoder layer in the last stage,
             # Only happens with block spec (TransformerBlockSubmodules) when using MoE.
-            spec = default_layer_spec(config)
+            # Re-derive all decoder layer specs and use the last one to get the correct
+            # layer type (dense vs MoE) for the MTP transformer layer.
+            from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_layer_specs
+
+            decoder_layer_specs = get_gpt_decoder_layer_specs(
+                config,
+                use_transformer_engine=True,
+                normalization=config.normalization,
+                qk_l2_norm=config.qk_l2_norm,
+            )
+            spec = decoder_layer_specs[-1]
         return get_gpt_mtp_block_spec(config, spec, use_transformer_engine=True, vp_stage=vp_stage)
     else:
         return None

tests/unit_tests/models/test_gpt_provider.py

@@ -374,6 +374,106 @@ def test_default_layer_spec_default_case(self, mock_te_full_spec, mock_te_spec):
         mock_te_spec.assert_called_once_with(provider)
         assert result == "te_spec"
 
+    def test_mtp_block_spec_returns_none_when_mtp_disabled(self):
+        """mtp_block_spec returns None when mtp_num_layers is unset."""
+        from megatron.bridge.models.gpt_provider import mtp_block_spec
+
+        provider = GPTModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+        )
+
+        assert mtp_block_spec(provider) is None
+
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_mtp_block_spec")
+    def test_mtp_block_spec_uses_callable_spec_directly_when_layer_specs_nonempty(self, mock_get_mtp):
+        """When the callable spec returns a non-empty block spec, use it as-is."""
+        from megatron.bridge.models.gpt_provider import mtp_block_spec
+
+        provider = GPTModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+            mtp_num_layers=1,
+        )
+
+        block_spec = Mock()
+        block_spec.layer_specs = ["layer_a", "layer_b"]
+        provider.transformer_layer_spec = lambda config: block_spec
+
+        mock_get_mtp.return_value = "mtp_spec"
+
+        result = mtp_block_spec(provider, vp_stage=None)
+
+        mock_get_mtp.assert_called_once_with(provider, block_spec, use_transformer_engine=True, vp_stage=None)
+        assert result == "mtp_spec"
+
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_decoder_layer_specs")
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_mtp_block_spec")
+    def test_mtp_block_spec_re_derives_last_decoder_spec_when_layer_specs_empty(
+        self, mock_get_mtp, mock_get_decoder_specs
+    ):
+        """When the last-stage spec has empty layer_specs (MoE block spec on the last PP stage),
+        re-derive all decoder layer specs and pass the last one to get_gpt_mtp_block_spec."""
+        from megatron.bridge.models.gpt_provider import mtp_block_spec
+
+        provider = GPTModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+            mtp_num_layers=1,
+        )
+
+        empty_block_spec = Mock()
+        empty_block_spec.layer_specs = []
+        provider.transformer_layer_spec = lambda config: empty_block_spec
+
+        dense_layer_spec = Mock(name="dense_layer_spec")
+        moe_layer_spec = Mock(name="moe_layer_spec")
+        mock_get_decoder_specs.return_value = [dense_layer_spec, moe_layer_spec]
+        mock_get_mtp.return_value = "mtp_spec"
+
+        result = mtp_block_spec(provider, vp_stage=2)
+
+        mock_get_decoder_specs.assert_called_once_with(
+            provider,
+            use_transformer_engine=True,
+            normalization=provider.normalization,
+            qk_l2_norm=provider.qk_l2_norm,
+        )
+        mock_get_mtp.assert_called_once_with(provider, moe_layer_spec, use_transformer_engine=True, vp_stage=2)
+        assert result == "mtp_spec"
+
+    @patch("megatron.core.models.gpt.gpt_layer_specs.get_gpt_mtp_block_spec")
+    def test_mtp_block_spec_passes_vp_stage_to_callable_spec(self, mock_get_mtp):
+        """When the transformer_layer_spec callable accepts vp_stage, it is forwarded."""
+        from megatron.bridge.models.gpt_provider import mtp_block_spec
+
+        provider = GPTModelProvider(
+            num_layers=2,
+            hidden_size=128,
+            num_attention_heads=4,
+            mtp_num_layers=1,
+        )
+
+        block_spec = Mock()
+        block_spec.layer_specs = ["layer_a"]
+        received_vp_stage = {}
+
+        def spec_fn(config, vp_stage=None):
+            received_vp_stage["vp_stage"] = vp_stage
+            return block_spec
+
+        provider.transformer_layer_spec = spec_fn
+        mock_get_mtp.return_value = "mtp_spec"
+
+        result = mtp_block_spec(provider, vp_stage=3)
+
+        assert received_vp_stage["vp_stage"] == 3
+        mock_get_mtp.assert_called_once_with(provider, block_spec, use_transformer_engine=True, vp_stage=3)
+        assert result == "mtp_spec"
+
     def test_dense_grouped_gemm_defaults_to_false(self):
         """GPTModelProvider.dense_grouped_gemm defaults to False."""
         provider = GPTModelProvider(