feat(megatron): support MTP in split training

nemo_rl/models/policy/tq_policy.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -65,6 +65,8 @@ def _aggregate_train_results(results: list[dict[str, Any]]) -> dict[str, Any]:
     }
     if "moe_metrics" in results[0]:
         out["moe_metrics"] = results[0]["moe_metrics"]
+    if "mtp_metrics" in results[0]:
+        out["mtp_metrics"] = results[0]["mtp_metrics"]
     all_mb_metrics: dict[str, list[Any]] = defaultdict(list)
     for r in results:
         for k, v in r["all_mb_metrics"].items():

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -1239,10 +1239,40 @@ def _split_step_state_init(
         if not isinstance(metric_normalizations, dict):
             metric_normalizations = {}
 
+        model_config = self._get_model_config()
+        mtp_num_layers = getattr(model_config, "mtp_num_layers", None)
+        mtp_enabled = mtp_num_layers is not None and mtp_num_layers > 0
+        mtp_detach_heads = bool(getattr(model_config, "mtp_detach_heads", False))
+        mtp_loss_scaling_factor = getattr(model_config, "mtp_loss_scaling_factor", 0.1)
+        loss_type = getattr(loss_fn, "loss_type", LossType.TOKEN_LEVEL)
+
+        # MTP is a token-summed auxiliary loss and therefore always needs the
+        # valid-token denominator. Under a sequence-level main loss, the split
+        # path can apply that distinct denominator only when mcore has isolated
+        # and tagged the detached MTP parameters. Attached-head gradients are
+        # already mixed into the backbone and cannot be corrected after the
+        # global counts become available at finish.
+        if (
+            mtp_enabled
+            and mtp_loss_scaling_factor != 0
+            and loss_type != LossType.TOKEN_LEVEL
+            and not mtp_detach_heads
+        ):
+            raise ValueError(
+                "MTP with a nonzero loss weight and sequence-level loss requires "
+                "policy.megatron_cfg.mtp_detach_heads=True on the SingleController "
+                "split training path because the MTP auxiliary gradient must be "
+                "normalized by valid tokens independently of the main loss. "
+                f"Got loss_type={loss_type}, mtp_num_layers={mtp_num_layers}, "
+                f"mtp_loss_scaling_factor={mtp_loss_scaling_factor}."
+            )
+
         return {
             "loss_fn": loss_fn,
-            "loss_type": getattr(loss_fn, "loss_type", LossType.TOKEN_LEVEL),
+            "loss_type": loss_type,
             "metric_normalizations": metric_normalizations,
+            "mtp_enabled": mtp_enabled,
+            "mtp_detach_heads": mtp_detach_heads,
             "gbs": gbs or self.cfg["train_global_batch_size"],
             "mbs": mbs or self.cfg["train_micro_batch_size"],
             "local_valid_seqs": torch.zeros((), dtype=torch.float64, device="cuda"),
@@ -1349,6 +1379,12 @@ def begin_train_step(
 
         state = self._split_step_state_init(loss_fn=loss_fn, gbs=gbs, mbs=mbs)
 
+        # Leave this unset so mcore falls back to config.grad_scale_func and
+        # inherits the optimizer's dynamic loss scale (especially for fp16).
+        # Also clear any transient callable left by an interrupted older step.
+        if state["mtp_enabled"]:
+            self._set_mtp_grad_scale_func(None)
+
         # Null the three mcore hooks that would fire a mid-step DP reduce:
         #   grad_sync_func — PP scheduler's direct call on last-MB boundaries
         #                    (PP>1 path).
@@ -1420,6 +1456,12 @@ def train_microbatch(
             # gradient finalization at all. Restore here; the caller is still
             # expected to invoke abort_train_step (idempotent on the saved
             # values) to drop ``_train_step_state``.
+            try:
+                self._set_mtp_grad_scale_func(None)
+            except Exception:
+                log.exception(
+                    "failed to clear MTP gradient scaling after train_microbatch error"
+                )
             try:
                 self._restore_saved_mcore_hooks(state)
             except Exception:
@@ -1455,6 +1497,15 @@ def _train_microbatch_body(
         state["local_valid_seqs"] = state["local_valid_seqs"] + call_local_seqs
         state["local_valid_toks"] = state["local_valid_toks"] + call_local_toks
 
+        # Match the synchronous Megatron path: derive the mask on the worker
+        # immediately before microbatch processing so sequence packing applies
+        # the same layout transformation to tokens and the MTP loss mask. The
+        # mask is worker-local derived data; it does not need a TQ schema field.
+        if state["mtp_enabled"] and "token_mask" in data and "sample_mask" in data:
+            data["mtp_loss_mask"] = data["token_mask"] * data["sample_mask"].unsqueeze(
+                -1
+            )
+
         # The number of chunks per optimizer step is a first-class property of
         # this path — it decides how many times gradients are accumulated before
         # a single reduce — but it was previously only recoverable by calibrating
@@ -1581,6 +1632,12 @@ def finish_train_step(self) -> dict[str, Any]:
             # the body got). Restore unconditionally so future steps run with
             # the right config. Leave ``_train_step_state`` for the caller's
             # abort_train_step to clear.
+            try:
+                self._set_mtp_grad_scale_func(None)
+            except Exception:
+                log.exception(
+                    "failed to clear MTP gradient scaling after finish_train_step error"
+                )
             try:
                 self._restore_saved_mcore_hooks(state)
             except Exception:
@@ -1614,6 +1671,17 @@ def _finish_train_step_body(self, state: dict[str, Any]) -> dict[str, Any]:
         # global mean grad; for reduce_scatter (dist-opt) it's the shard.
         # Either way, opt.step sees the right-normalized gradient.
         self.model.scale_gradients(inv_n)
+        # The uniform rescale gives MTP the main loss's denominator. Correct
+        # detached, MTP-tagged parameters back to the valid-token denominator
+        # used by the synchronous path. For token-level loss the factor is 1.
+        # This runs before gradient reduction; scaling and reduction are linear.
+        if state["mtp_enabled"] and state["mtp_detach_heads"]:
+            self._scale_mtp_param_grads(
+                float((n_safe / global_valid_toks.clamp(min=1)).item())
+            )
+        # No more forward/backward calls remain in this step. Clear the
+        # callable before optimizer/scheduler/checkpoint state can serialize it.
+        self._set_mtp_grad_scale_func(None)
 
         # End-of-step gradient finalization, exactly once per optimizer step.
         # ``begin_train_step`` nulled ``finalize_model_grads_func`` so mcore's
@@ -1668,6 +1736,11 @@ def _finish_train_step_body(self, state: dict[str, Any]) -> dict[str, Any]:
         # opt.step clips internally (clip_grad config); operates on the
         # already-rescaled grad. Returns (success, grad_norm, num_zeros).
         update_successful, grad_norm, num_zeros_in_grad = self.optimizer.step()
+        mtp_grad_norm = (
+            self.optimizer.grad_norms_by_group.get("mtp")
+            if state["mtp_enabled"]
+            else None
+        )
 
         pg_collection = get_pg_collection(self.model)
         update_successful = logical_and_across_model_parallel_group(
@@ -1679,6 +1752,13 @@ def _finish_train_step_body(self, state: dict[str, Any]) -> dict[str, Any]:
         num_zeros_in_grad = reduce_max_stat_across_model_parallel_group(
             num_zeros_in_grad, mp_group=pg_collection.mp
         )
+        if state["mtp_enabled"]:
+            # MTP parameters live on the last PP stage. Make their independently
+            # clipped grad norm visible to every model-parallel rank before the
+            # driver selects a replica leader's result.
+            mtp_grad_norm = reduce_max_stat_across_model_parallel_group(
+                mtp_grad_norm, mp_group=pg_collection.mp
+            )
 
         # Mirrors train(): without re-enabling the pre-hook __init__ removed, the
         # param all-gather never runs and each forward sees only its own shard.
@@ -1831,6 +1911,12 @@ def _scale_metric(name: str, value: Any) -> Any:
             if moe_metrics:
                 metrics["moe_metrics"] = moe_metrics
 
+        self._collect_mtp_metrics(
+            metrics,
+            state["total_num_microbatches"],
+            mtp_grad_norm,
+        )
+
         self._train_step_state = None
         return metrics
 
@@ -1839,12 +1925,15 @@ def abort_train_step(self) -> None:
         state = getattr(self, "_train_step_state", None)
         if state is None:
             return
-        # Restore the mcore hooks first so the model is back to a normal
-        # state before zero_grad_buffer touches anything.
-        self._restore_saved_mcore_hooks(state)
-        self.model.zero_grad_buffer()
-        self.optimizer.zero_grad()
-        self._train_step_state = None
+        # Drop the step-local MTP scaler and restore the mcore hooks before
+        # zero_grad_buffer touches anything.
+        try:
+            self._set_mtp_grad_scale_func(None)
+        finally:
+            self._restore_saved_mcore_hooks(state)
+            self.model.zero_grad_buffer()
+            self.optimizer.zero_grad()
+            self._train_step_state = None
 
     @wrap_with_nvtx_name("megatron_policy_worker/get_logprobs")
     def get_logprobs(
@@ -2226,6 +2315,26 @@ def _set_mtp_grad_scale_func(self, func):
         if config is not None:
             config.mtp_grad_scale_func = func
 
+    def _scale_mtp_param_grads(self, factor: float) -> None:
+        """Scale detached MTP parameters' gradients by ``factor``.
+
+        MCore tags every MTP parameter ``grad_norm_group='mtp'`` when
+        ``mtp_detach_heads`` is enabled. In that configuration the auxiliary
+        loss reaches no shared parameters, so its denominator can be corrected
+        independently. ``main_grad`` is a view into the DDP gradient buffer.
+
+        Args:
+            factor: Multiplier for MTP gradients. ``1.0`` is a no-op.
+        """
+        if factor == 1.0:
+            return
+        for param in self.model.parameters():
+            if getattr(param, "grad_norm_group", None) != "mtp":
+                continue
+            main_grad = getattr(param, "main_grad", None)
+            if main_grad is not None:
+                main_grad.mul_(factor)
+
     def _get_model_config(self):
         """Get the underlying model config (handle Float16Module wrapper)."""
         model = self.model