[Dev] Cherry-pick MTP detach heads

megatron/core/transformer/multi_token_prediction.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.
 from __future__ import annotations
 
 import warnings
@@ -881,7 +881,7 @@ def set_loss_scale(scale: torch.Tensor):
 
 def process_mtp_loss(
     hidden_states: Tensor,
-    labels: Optional[Tensor],
+    labels: Tensor,
     loss_mask: Optional[Tensor],
     output_layer: Callable,
     output_weight: Optional[Tensor],
@@ -937,6 +937,12 @@ def process_mtp_loss(
         )
         derived_labels_from_input_ids = True
 
+    if config.mtp_detach_heads:
+        if output_weight is not None:
+            output_weight = output_weight.detach()
+        else:
+            output_weight = output_layer.weight.detach()
+
     mtp_labels = labels.clone()
     if loss_mask is None:
         loss_mask = torch.ones_like(mtp_labels)
@@ -948,23 +954,6 @@ def process_mtp_loss(
             loss_mask, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
         )
 
-    output_weight_for_mtp = output_weight
-    output_layer_for_mtp = output_layer
-    if config.mtp_isolated_loss:
-        if output_weight_for_mtp is not None:
-            output_weight_for_mtp = output_weight_for_mtp.detach()
-        if isinstance(output_layer, torch.nn.Module):
-            output_layer_params = {
-                name: param.detach() for name, param in output_layer.named_parameters()
-            }
-            output_layer_buffers = dict(output_layer.named_buffers())
-            output_layer_state = {**output_layer_params, **output_layer_buffers}
-
-            def output_layer_for_mtp(input_: Tensor, **kwargs):
-                return torch.func.functional_call(
-                    output_layer, output_layer_state, args=(input_,), kwargs=kwargs
-                )
-
     # Store the original number of tokens before rolling for proper normalization
     # when calculate_per_token_loss is enabled. This ensures MTP gradients are
     # correctly scaled relative to the main loss gradients in finalize_model_grads.
@@ -981,9 +970,9 @@ def output_layer_for_mtp(input_: Tensor, **kwargs):
             loss_mask, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params
         )
         if fuse_linear_cross_entropy:
-            mtp_loss = output_layer_for_mtp(
+            mtp_loss = output_layer(
                 hidden_states_list[mtp_layer_number + 1],
-                weight=output_weight_for_mtp,
+                weight=output_weight,
                 runtime_gather_output=runtime_gather_output,
                 output_cross_entropy_loss=True,
                 labels=mtp_labels,
@@ -992,9 +981,9 @@ def output_layer_for_mtp(input_: Tensor, **kwargs):
             # acceptance counts cannot be computed for this layer.
             mtp_logits = None
         else:
-            mtp_logits, _ = output_layer_for_mtp(
+            mtp_logits, _ = output_layer(
                 hidden_states_list[mtp_layer_number + 1],
-                weight=output_weight_for_mtp,
+                weight=output_weight,
                 runtime_gather_output=runtime_gather_output,
             )
             if scale_logits_fn is not None:
@@ -1296,10 +1285,18 @@ def _get_embeddings(
             )
         # embedding
         decoder_input = embedding(input_ids=input_ids, position_ids=position_ids)
-        if self.config.mtp_isolated_loss:
+
+        if self.config.mtp_detach_heads:
             decoder_input = decoder_input.detach()
 
         hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
+        # make_viewless_tensor no-ops when hidden_states is not a view (_base is None),
+        # which happens after detach() with mtp_detach_heads. Activation
+        # checkpointing (CheckpointFunction.apply) requires at least one input tensor
+        # with requires_grad=True to produce a differentiable output, so we ensure it
+        # here to maintain gradient flow to MTP layer parameters.
+        if not hidden_states.requires_grad:
+            hidden_states.requires_grad_(True)
 
         return input_ids, position_ids, padding_mask, decoder_input, hidden_states
 
@@ -2052,11 +2049,10 @@ def forward(
             hidden_states = mhc_chunks[offset]
         else:
             hidden_states = hidden_states_list[offset]
-        if self.config.mtp_isolated_loss:
-            hidden_states = hidden_states.detach().requires_grad_(True)
-            hidden_states = make_viewless_tensor(
-                inp=hidden_states, requires_grad=True, keep_graph=False
-            )
+
+        if self.config.mtp_detach_heads:
+            hidden_states = hidden_states.detach()
+
         for iteration in range(self.config.mtp_num_layers):
             layer_idx = 0 if self.mtp_use_repeated_layer else iteration
             (hidden_states, input_ids, position_ids, padding_mask) = self.layers[layer_idx](

megatron/core/transformer/transformer_config.py

@@ -81,18 +81,14 @@ class TransformerConfig(ModelParallelConfig):
     which serves as an additional training objective.
     """
 
-    mtp_isolated_loss: bool = False
-    """If True, MTP loss only updates MTP module parameters. The MTP loss graph is
-    detached from the main decoder, shared embeddings, and output layer weights.
-
-    For online RL, keep ``labels=None`` so the main LM head returns logits for the
-    external RL loss. MTP auxiliary loss can still be trained by deriving its labels
-    from ``input_ids`` in the MTP loss path; this option isolates that auxiliary loss
-    from the main model parameters."""
-
     mtp_use_repeated_layer: bool = False
     """Use a single MTP layer repeatedly instead of multiple separate layers."""
 
+    mtp_detach_heads: bool = False
+    """If True, detach MTP head inputs from the main model graph.
+    This prevents MTP loss gradients from flowing back to the main model,
+    only training the MTP heads themselves."""
+
     mtp_hybrid_override_pattern: Optional[str] = None
     """DEPRECATED: Use unified hybrid_layer_pattern instead.
     Legacy argument for loading old checkpoints.

tests/unit_tests/models/test_hybrid_moe_model.py

@@ -230,8 +230,8 @@
     "mup_embedding_mult": 1.0,
     "mup_output_mult": 1.0,
     "mup_width_mult": 1.0,
+    "mtp_detach_heads": False,
     "mtp_hybrid_override_pattern": None,
-    "mtp_isolated_loss": False,
     "mtp_loss_scaling_factor": 0.1,
     "mtp_num_layers": None,
     "mtp_standalone": False,