Add HybridModel MTP runtime controls

megatron/core/models/hybrid/hybrid_model.py

@@ -426,12 +426,19 @@ def forward(
         loss_mask: Optional[Tensor] = None,
         packed_seq_params: Optional[PackedSeqParams] = None,
         padding_mask: Optional[Tensor] = None,
+        run_mtp_forward: bool = True,
     ) -> Tensor:
         """Forward function of the Hybrid model. This function passes the input tensors
         through the embedding layer, and then the decoder and finally into the post
         processing layer (optional).
 
         It either returns the Loss values if labels are given or the final hidden units
+
+        Args:
+            run_mtp_forward (bool): Whether to execute the non-inference MTP forward and attach
+                its auxiliary loss. Disabling this leaves the MTP module and checkpoint state
+                intact. Serial MTP speculative decoding is configured through the inference
+                context and controller. Defaults to True.
         """
         # If decoder_input is provided (not None), then input_ids and position_ids are ignored.
         # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input.
@@ -535,7 +542,9 @@ def forward(
             and inference_context.num_speculative_tokens > 0
         )
 
-        mtp_forward_ran = self.mtp_process and not (in_inference_mode or is_spec_decode)
+        mtp_forward_ran = (
+            self.mtp_process and run_mtp_forward and not (in_inference_mode or is_spec_decode)
+        )
         if mtp_forward_ran:
             hidden_states = self.mtp(
                 input_ids=input_ids,
@@ -567,7 +576,7 @@ def forward(
                     # Non-block scope: direct assignment; the controller will set
                     # this back to None after reading to allow GC.
                     inference_context.mtp_decoder_hidden_states = hidden_states
-            elif not in_inference_mode:
+            elif mtp_forward_ran:
                 # For RL (labels is None), process_mtp_loss derives labels from
                 # input_ids to match the SFT label format.
                 hidden_states = process_mtp_loss(

tests/unit_tests/transformer/test_multi_token_prediction.py

@@ -9,6 +9,7 @@
 
 from megatron.core.enums import ModelType
 from megatron.core.extensions.transformer_engine import HAVE_TE
+from megatron.core.inference.utils import InferenceMode
 from megatron.core.models.gpt.gpt_layer_specs import (
     get_gpt_layer_local_spec,
     get_gpt_layer_with_transformer_engine_spec,
@@ -1370,6 +1371,159 @@ def get_batch(self, seq_length, micro_batch_size):
         }
         return batch
 
+    @staticmethod
+    def _make_forward_stub(mtp_process=True, post_process=True):
+        hidden_states = torch.arange(4, dtype=torch.float32).reshape(2, 1, 2)
+        call_counts = {"mtp": 0, "loss": 0}
+
+        def decoder(**kwargs):
+            hidden_states = kwargs["hidden_states"]
+            return hidden_states.unwrap() if hasattr(hidden_states, "unwrap") else hidden_states
+
+        def mtp(**kwargs):
+            call_counts["mtp"] += 1
+            decoder_hidden_states = kwargs["hidden_states"]
+            return torch.cat((decoder_hidden_states, decoder_hidden_states + 100.0), dim=0)
+
+        def output_layer(output, **kwargs):
+            return output, None
+
+        model_attrs = dict(
+            config=types.SimpleNamespace(
+                fine_grained_activation_offloading=False,
+                moe_paged_stash=False,
+                mtp_num_layers=1,
+                use_mup=False,
+                inference_cuda_graph_scope=None,
+            ),
+            pre_process=False,
+            post_process=post_process,
+            position_embedding_type='none',
+            decoder=decoder,
+            share_embeddings_and_output_weights=False,
+            mtp_process=mtp_process,
+            embedding=None,
+            output_layer=output_layer,
+            training=True,
+            compute_language_model_loss=lambda labels, logits: logits,
+            pg_collection=types.SimpleNamespace(cp=None),
+            tp_group=None,
+            _scale_logits=lambda logits: logits,
+        )
+        if mtp_process:
+            model_attrs["mtp"] = mtp
+        model = types.SimpleNamespace(**model_attrs)
+        return model, hidden_states, call_counts
+
+    @pytest.mark.parametrize(
+        ("forward_kwargs", "expected_mtp_calls", "expected_loss_calls"),
+        [
+            pytest.param({}, 1, 1, id="defaults"),
+            pytest.param({"run_mtp_forward": True}, 1, 1, id="explicitly-enabled"),
+            pytest.param({"run_mtp_forward": False}, 0, 0, id="disabled"),
+        ],
+    )
+    def test_forward_mtp_control(
+        self, monkeypatch, forward_kwargs, expected_mtp_calls, expected_loss_calls
+    ):
+        """Test enabled and disabled HybridModel MTP forward execution and loss processing."""
+        model, hidden_states, call_counts = self._make_forward_stub()
+
+        def process_mtp_loss_spy(**kwargs):
+            call_counts["loss"] += 1
+            return torch.chunk(kwargs["hidden_states"], 1 + kwargs["config"].mtp_num_layers, dim=0)[
+                0
+            ]
+
+        monkeypatch.setattr(
+            "megatron.core.models.hybrid.hybrid_model.process_mtp_loss", process_mtp_loss_spy
+        )
+
+        output = HybridModel.forward(
+            model,
+            input_ids=torch.zeros(1, 2, dtype=torch.long),
+            position_ids=torch.arange(2).unsqueeze(0),
+            attention_mask=None,
+            decoder_input=hidden_states,
+            **forward_kwargs,
+        )
+
+        torch.testing.assert_close(output, hidden_states.transpose(0, 1).contiguous())
+        assert call_counts == {"mtp": expected_mtp_calls, "loss": expected_loss_calls}
+
+    @pytest.mark.parametrize(
+        "forward_kwargs",
+        [
+            pytest.param({"run_mtp_forward": True}, id="enabled"),
+            pytest.param({"run_mtp_forward": False}, id="disabled"),
+        ],
+    )
+    def test_forward_mtp_control_on_non_mtp_rank(self, forward_kwargs):
+        """Test that the MTP control does not access MTP on a pipeline rank without it."""
+        model, hidden_states, call_counts = self._make_forward_stub(mtp_process=False)
+
+        output = HybridModel.forward(
+            model,
+            input_ids=torch.zeros(1, 2, dtype=torch.long),
+            position_ids=torch.arange(2).unsqueeze(0),
+            attention_mask=None,
+            decoder_input=hidden_states,
+            **forward_kwargs,
+        )
+
+        torch.testing.assert_close(output, hidden_states.transpose(0, 1).contiguous())
+        assert call_counts == {"mtp": 0, "loss": 0}
+
+    def test_forward_mtp_disabled_before_post_process(self):
+        """Test that disabled MTP is skipped before the pipeline boundary."""
+        model, hidden_states, call_counts = self._make_forward_stub(post_process=False)
+
+        output = HybridModel.forward(
+            model,
+            input_ids=torch.zeros(1, 2, dtype=torch.long),
+            position_ids=torch.arange(2).unsqueeze(0),
+            attention_mask=None,
+            decoder_input=hidden_states,
+            run_mtp_forward=False,
+        )
+
+        torch.testing.assert_close(output, hidden_states)
+        assert call_counts == {"mtp": 0, "loss": 0}
+
+    def test_forward_mtp_control_does_not_disable_speculative_decoding(self, monkeypatch):
+        """Test that the MTP control does not disable speculative decoding state."""
+        model, hidden_states, call_counts = self._make_forward_stub()
+        inference_context = types.SimpleNamespace(
+            is_dynamic_batching=lambda: True,
+            num_speculative_tokens=1,
+            mtp_decoder_hidden_states=None,
+            config=types.SimpleNamespace(materialize_only_last_token_logits=False),
+        )
+
+        def unexpected_process_mtp_loss(**kwargs):
+            call_counts["loss"] += 1
+            raise AssertionError("MTP loss processing must not run during speculative decoding")
+
+        monkeypatch.setattr(
+            "megatron.core.models.hybrid.hybrid_model.process_mtp_loss", unexpected_process_mtp_loss
+        )
+
+        with InferenceMode.active():
+            output = HybridModel.forward(
+                model,
+                input_ids=torch.zeros(1, 2, dtype=torch.long),
+                position_ids=torch.arange(2).unsqueeze(0),
+                attention_mask=None,
+                decoder_input=hidden_states,
+                inference_context=inference_context,
+                runtime_gather_output=True,
+                run_mtp_forward=False,
+            )
+
+        torch.testing.assert_close(output, hidden_states.transpose(0, 1).contiguous())
+        torch.testing.assert_close(inference_context.mtp_decoder_hidden_states, hidden_states)
+        assert call_counts == {"mtp": 0, "loss": 0}
+
     @pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available")
     @pytest.mark.parametrize(("tp", "cp"), [(1, 1), (2, 1)])
     def test_sharded_state_dict_mamba(self, tp, cp):