feat: add MTP (multi-token prediction) training support for nemotron

examples/configs/grpo_math_1B.yaml

@@ -181,6 +181,14 @@ policy:
     moe_enable_deepep: false
     moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
+    # Multi-Token Prediction (MTP). mtp_num_layers=0 disables MTP.
+    mtp_num_layers: 0
+    # MTP loss weight added to the main next-token loss (0.0 disables the MTP loss contribution).
+    mtp_loss_scaling_factor: 0.0
+    # When true, repeat one MTP layer mtp_num_layers times instead of using distinct layers.
+    mtp_use_repeated_layer: false
+    # When true, detach MTP heads so MTP loss does not affect main-model gradients.
+    mtp_detach_heads: false
     gradient_accumulation_fusion: false
 
     peft:

nemo_rl/algorithms/grpo.py

@@ -2389,6 +2389,10 @@ def grpo_train(
                     metrics.update(
                         {f"moe/{k}": v for k, v in train_results["moe_metrics"].items()}
                     )
+                if "mtp_metrics" in train_results:
+                    metrics.update(
+                        {f"mtp/{k}": v for k, v in train_results["mtp_metrics"].items()}
+                    )
                 if master_config.grpo["use_dynamic_sampling"]:
                     metrics["filtered_reward"] = rewards.numpy()
                     metrics["reward"] = repeated_batch["total_reward"].numpy()
@@ -3593,6 +3597,10 @@ def async_grpo_train(
                     metrics.update(
                         {f"moe/{k}": v for k, v in train_results["moe_metrics"].items()}
                     )
+                if "mtp_metrics" in train_results:
+                    metrics.update(
+                        {f"mtp/{k}": v for k, v in train_results["mtp_metrics"].items()}
+                    )
                 metrics.update(train_results["all_mb_metrics"])
                 for k, v in metrics.items():
                     if k in {"probs_ratio_min", "probs_ratio_clamped_min"}:

nemo_rl/models/megatron/common.py

@@ -21,6 +21,7 @@
     get_moe_layer_wise_logging_tracker,
     reduce_aux_losses_tracker_across_ranks,
 )
+from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper
 
 
 def _round_up_to_multiple(value: int, multiple: int) -> int:
@@ -164,3 +165,34 @@ def get_moe_metrics(
 
     clear_aux_losses_tracker()
     return metrics
+
+
+def get_mtp_metrics() -> dict[str, Any]:
+    """Returns Multi-Token Prediction (MTP) loss and acceptance rate metrics.
+
+    This function reduces MTP metrics across ranks and returns a dictionary of metrics.
+
+    Returns:
+        dict[str, Any]: A flat dict of metrics. Each MTP layer's loss is returned
+        under the key "mtp_{i}_loss" and acceptance rate under "mtp_{i}_acceptance_rate"
+        where i is 1-indexed (matching Megatron-LM).
+    """
+    MTPLossLoggingHelper.reduce_metrics_in_tracker()
+    tracker = MTPLossLoggingHelper.tracker
+
+    metrics: dict[str, Any] = {}
+    if "loss_values" in tracker:
+        mtp_losses = tracker["loss_values"].float()
+        mtp_corrects = tracker.get("correct_values", torch.zeros_like(mtp_losses))
+        mtp_totals = tracker.get("total_values", torch.ones_like(mtp_losses))
+        mtp_num_layers = mtp_losses.shape[0]
+
+        # Log per-layer losses and acceptance rates
+        for i in range(mtp_num_layers):
+            metrics[f"mtp_{i + 1}_loss"] = float(mtp_losses[i].item())
+            # Compute acceptance rate as percentage
+            acceptance_rate = (mtp_corrects[i] / max(mtp_totals[i], 1)) * 100.0
+            metrics[f"mtp_{i + 1}_acceptance_rate"] = float(acceptance_rate.item())
+
+        MTPLossLoggingHelper.clean_metrics_in_tracker()
+    return metrics

nemo_rl/models/megatron/data.py

@@ -40,6 +40,7 @@ class ProcessedInputs:
     position_ids: Optional[torch.Tensor]
     packed_seq_params: Optional[PackedSeqParams]
     cu_seqlens_padded: Optional[torch.Tensor]
+    mtp_loss_mask: Optional[torch.Tensor] = None
 
 
 @dataclass
@@ -66,6 +67,7 @@ class ProcessedMicrobatch:
     position_ids: Optional[torch.Tensor]
     packed_seq_params: Optional[PackedSeqParams]
     cu_seqlens_padded: Optional[torch.Tensor]
+    mtp_loss_mask: Optional[torch.Tensor] = None
 
 
 def make_processed_microbatch_iterator(
@@ -121,6 +123,7 @@ def make_processed_microbatch_iterator(
             position_ids=processed_inputs.position_ids,
             packed_seq_params=processed_inputs.packed_seq_params,
             cu_seqlens_padded=processed_inputs.cu_seqlens_padded,
+            mtp_loss_mask=processed_inputs.mtp_loss_mask,
         )
 
 
@@ -251,6 +254,7 @@ def process_microbatch(
         seq_lengths = None  # Will be set if using packed sequences
         cu_seqlens = None
         cu_seqlens_padded = None
+        mtp_loss_mask = None
 
         if pack_sequences:
             # For packed sequences with padded input, we need sequence lengths
@@ -302,6 +306,24 @@ def process_microbatch(
                     cp_size=get_context_parallel_world_size(),
                 )
 
+                # Pack pre-computed mtp_loss_mask the same way as input_ids
+                if "mtp_loss_mask" in data_dict:
+                    (
+                        _,
+                        mtp_loss_mask,
+                        _,
+                        _,
+                        _,
+                    ) = _pack_sequences_for_megatron(
+                        data_dict["mtp_loss_mask"],
+                        seq_lengths,
+                        pad_individual_seqs_to_multiple_of,
+                        pad_packed_seq_to_multiple_of,
+                        pad_full_seq_to,
+                        cp_rank=get_context_parallel_rank(),
+                        cp_size=get_context_parallel_world_size(),
+                    )
+
                 # For packed sequences, position_ids and attention_mask are typically None
                 # The PackedSeqParams handles all necessary sequence information
                 position_ids = None
@@ -317,13 +339,16 @@ def process_microbatch(
                 eod_mask_loss=False,
                 pad_mask_loss=False,
             )
+            if "mtp_loss_mask" in data_dict:
+                mtp_loss_mask = data_dict["mtp_loss_mask"]
     return ProcessedInputs(
         input_ids=input_ids,
         input_ids_cp_sharded=input_ids_cp_sharded,
         attention_mask=attention_mask,
         position_ids=position_ids,
         packed_seq_params=packed_seq_params,
         cu_seqlens_padded=cu_seqlens_padded,
+        mtp_loss_mask=mtp_loss_mask,
     )
 
 

nemo_rl/models/megatron/setup.py

@@ -660,8 +660,19 @@ def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None:
 
 
 def _apply_mtp_config(model_cfg: Any, config: PolicyConfig) -> None:
-    if "mtp_num_layers" in config["megatron_cfg"]:
-        model_cfg.mtp_num_layers = config["megatron_cfg"]["mtp_num_layers"]
+    """Apply Multi-Token Prediction settings onto the mcore model config."""
+    megatron_cfg = config["megatron_cfg"]
+    if "mtp_num_layers" in megatron_cfg:
+        # In mcore, mtp_num_layers is both the number of MTP layers (when
+        # mtp_use_repeated_layer is False) and the number of times the MTP layer
+        # is repeated (when mtp_use_repeated_layer is True).
+        model_cfg.mtp_num_layers = megatron_cfg["mtp_num_layers"]
+    if "mtp_loss_scaling_factor" in megatron_cfg:
+        model_cfg.mtp_loss_scaling_factor = megatron_cfg["mtp_loss_scaling_factor"]
+    if "mtp_use_repeated_layer" in megatron_cfg:
+        model_cfg.mtp_use_repeated_layer = megatron_cfg["mtp_use_repeated_layer"]
+    if "mtp_detach_heads" in megatron_cfg:
+        model_cfg.mtp_detach_heads = megatron_cfg["mtp_detach_heads"]
 
 
 def _apply_precision_config(

nemo_rl/models/megatron/train.py

@@ -73,6 +73,7 @@ def model_forward(
     attention_mask: torch.Tensor,
     packed_seq_params: Optional[PackedSeqParams] = None,
     defer_fp32_logits: Optional[bool] = False,
+    mtp_loss_mask: Optional[torch.Tensor] = None,
     straggler_timer: Optional[StragglerDetector] = None,
     use_linear_ce_fusion_loss: bool = False,
 ) -> torch.Tensor:
@@ -86,6 +87,7 @@ def model_forward(
         attention_mask: Attention mask for the sequence
         packed_seq_params: Parameters for packed sequences (optional)
         defer_fp32_logits: Whether to skip the conversion of logits to fp32
+        mtp_loss_mask: MTP loss mask to exclude prompt tokens from MTP loss (optional)
         straggler_timer: Straggler detector for profiling the forward pass
         use_linear_ce_fusion_loss: Whether to use linear CE fusion loss
 
@@ -102,6 +104,11 @@ def model_forward(
     # Mamba models currently do not support packed_seq_params
     if packed_seq_params is not None:
         additional_kwargs["packed_seq_params"] = packed_seq_params
+
+    # Pass MTP loss mask to exclude prompt tokens from MTP loss
+    if mtp_loss_mask is not None:
+        additional_kwargs["loss_mask"] = mtp_loss_mask
+
     if defer_fp32_logits:
         additional_kwargs["fp32_output"] = False
     if use_linear_ce_fusion_loss:
@@ -186,6 +193,7 @@ def forward_with_post_processing_fn(
     position_ids = processed_mb.position_ids
     packed_seq_params = processed_mb.packed_seq_params
     cu_seqlens_padded = processed_mb.cu_seqlens_padded
+    mtp_loss_mask = processed_mb.mtp_loss_mask
 
     # Insert hook to capture hidden states and embeddings for draft model training if draft_model is provided
     capture_context, capture = get_capture_context(model, enable_hidden_capture)
@@ -198,6 +206,7 @@ def forward_with_post_processing_fn(
             attention_mask=attention_mask,
             packed_seq_params=packed_seq_params,
             defer_fp32_logits=defer_fp32_logits,
+            mtp_loss_mask=mtp_loss_mask,
             straggler_timer=straggler_timer,
             use_linear_ce_fusion_loss=use_linear_ce_fusion_loss,
         )

nemo_rl/models/policy/__init__.py

@@ -296,6 +296,12 @@ class MegatronConfig(TypedDict):
     linear_ce_fusion_chunk_size: NotRequired[int]
     # When mtp_num_layers=0, Multi-Token Prediction is disabled.
     mtp_num_layers: NotRequired[int]
+    # MTP loss weight added to the main next-token loss (0.0 disables the MTP loss contribution).
+    mtp_loss_scaling_factor: NotRequired[float]
+    # When True, repeat a single MTP layer mtp_num_layers times instead of using distinct layers.
+    mtp_use_repeated_layer: NotRequired[bool]
+    # When True, detach MTP heads from the main model so MTP loss does not affect main-model gradients.
+    mtp_detach_heads: NotRequired[bool]
     # When True, clear the RotaryEmbedding LRU cache and MoE token dispatcher
     # routing tensors in offload_before_refit (before weight transfer to the
     # inference engine). Useful when training and logprob runs use different

nemo_rl/models/policy/lm_policy.py

@@ -729,6 +729,8 @@ def train(
         }
         if "moe_metrics" in results[0]:
             aggregated_results["moe_metrics"] = results[0]["moe_metrics"]
+        if "mtp_metrics" in results[0]:
+            aggregated_results["mtp_metrics"] = results[0]["mtp_metrics"]
 
         if self.flops_tracker is not None:
             aggregated_results["total_flops"] = self.flops_tracker.total_flops

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -510,6 +510,14 @@ def train(
                 global_valid_seqs = gb_result["global_valid_seqs"]
                 global_valid_toks = gb_result["global_valid_toks"]
 
+                # Pre-compute MTP loss mask from token_mask and sample_mask
+                # before microbatch processing, so process_microbatch can pack it
+                if "token_mask" in batch and "sample_mask" in batch:
+                    mtp_loss_mask = batch["token_mask"] * batch[
+                        "sample_mask"
+                    ].unsqueeze(-1)
+                    batch["mtp_loss_mask"] = mtp_loss_mask
+
                 (
                     data_iterator,
                     num_microbatches,
@@ -546,6 +554,10 @@ def train(
                         self.optimizer.zero_grad()
                         self._copy_main_params_to_param_buffer()
 
+                    # Set mtp_grad_scale_func for MTP loss scaling (scales by valid tokens)
+                    mtp_scale = 1.0 / global_valid_toks.clamp(min=1).float()
+                    self._set_mtp_grad_scale_func(lambda: mtp_scale)
+
                     # Forward pass.
                     draft_enabled = "draft" in self.cfg and self.cfg["draft"]["enabled"]
                     losses_reduced = megatron_forward_backward(
@@ -568,6 +580,10 @@ def train(
                         ),
                     )
 
+                # Clear mtp_grad_scale_func after the forward-backward pass so
+                # it doesn't get serialized in the run_config.yaml when saving
+                self._set_mtp_grad_scale_func(None)
+
                 # Empty unused memory.
                 if self.cfg["megatron_cfg"]["empty_unused_memory_level"] >= 1:
                     torch.cuda.empty_cache()
@@ -688,6 +704,9 @@ def train(
             )
             if moe_metrics:
                 metrics["moe_metrics"] = moe_metrics
+        # Collect MTP metrics (kept out of train()'s body so cloudpickle does not
+        # pull an unpicklable torch ConfigModuleInstance into the worker actor).
+        self._collect_mtp_metrics(metrics)
         return metrics
 
     @wrap_with_nvtx_name("megatron_policy_worker/get_logprobs")
@@ -1207,6 +1226,36 @@ def prepare_refit_info(self) -> None:
 
         return refit_param_info_hf
 
+    def _collect_mtp_metrics(self, metrics: dict[str, Any]) -> None:
+        """Add Multi-Token Prediction metrics to ``metrics`` when MTP is enabled.
+
+        get_mtp_metrics is imported lazily (not a module global) so cloudpickle
+        does not pull an unpicklable torch ConfigModuleInstance into the worker
+        actor's serialization.
+        """
+        mtp_num_layers = getattr(self.model.config, "mtp_num_layers", None)
+        if mtp_num_layers is not None and mtp_num_layers > 0:
+            from nemo_rl.models.megatron.common import get_mtp_metrics
+
+            mtp_metrics = get_mtp_metrics()
+            if mtp_metrics:
+                metrics["mtp_metrics"] = mtp_metrics
+
+    def _set_mtp_grad_scale_func(self, func):
+        """Set mtp_grad_scale_func on the model config for MTP loss scaling."""
+        config = self._get_model_config()
+        if config is not None:
+            config.mtp_grad_scale_func = func
+
+    def _get_model_config(self):
+        """Get the underlying model config (handle Float16Module wrapper)."""
+        model = self.model
+        if hasattr(model, "module") and hasattr(model.module, "config"):
+            return model.module.config
+        elif hasattr(model, "config"):
+            return model.config
+        return None
+
     def _calculate_refit_param_info(self) -> list[tuple[str, int]]:
         """Calculate parameter information for refit.