fix(eagle3): isolate draft grad clipping

nemo_rl/algorithms/grpo.py

@@ -3192,6 +3192,10 @@ def grpo_train(
                     metrics.update(
                         {f"mtp/{k}": v for k, v in train_results["mtp_metrics"].items()}
                     )
+                if "draft_grad_norm" in train_results:
+                    metrics["draft_grad_norm"] = train_results[
+                        "draft_grad_norm"
+                    ].numpy()
                 if master_config.grpo["use_dynamic_sampling"]:
                     metrics["filtered_reward"] = rewards.numpy()
                     metrics["reward"] = repeated_batch["total_reward"].numpy()
@@ -4596,6 +4600,10 @@ def async_grpo_train(
                     metrics.update(
                         {f"mtp/{k}": v for k, v in train_results["mtp_metrics"].items()}
                     )
+                if "draft_grad_norm" in train_results:
+                    metrics["draft_grad_norm"] = train_results[
+                        "draft_grad_norm"
+                    ].numpy()
                 metrics.update(train_results["all_mb_metrics"])
                 metrics.update(penalty_metrics)
                 for k, v in metrics.items():

nemo_rl/algorithms/loss/loss_functions.py

@@ -92,6 +92,8 @@ def __call__(
                 False,
             )
         else:
+            # teacher_logits is already detached at the call site (utils.py);
+            # match DistributedCrossEntropy semantics.
             teacher_probs = torch.nn.functional.softmax(teacher_logits, dim=-1)
             student_log_probs = torch.nn.functional.log_softmax(student_logits, dim=-1)
             per_token_loss = -(teacher_probs * student_log_probs).sum(dim=-1)

nemo_rl/models/generation/vllm/vllm_worker.py

@@ -644,6 +644,22 @@ def stop_gpu_profiling(self) -> None:
         if self.llm is not None:
             self.llm.collective_rpc("stop_gpu_profiling", args=tuple())
 
+    @staticmethod
+    def _spec_decode_max_tokens(
+        base_max_tokens: int,
+        input_len: int,
+        max_model_len: int,
+        spec_lookahead: int,
+    ) -> int:
+        """Clamp max_tokens so speculative decoding never reads past max_model_len.
+
+        The drafter looks `spec_lookahead` tokens ahead, so generation must stop
+        at least `spec_lookahead + 1` tokens before the boundary.
+        """
+        return max(
+            1, min(base_max_tokens, max_model_len - input_len - (spec_lookahead + 1))
+        )
+
     @staticmethod
     def _patch_vllm_nsight_config() -> None:
         """Override vLLM's nsight config for internal TP workers to use deferred capture.
@@ -807,6 +823,26 @@ def generate(
             stop_strings=stop_strings,
         )
 
+        # vLLM 0.20 eagle3 spec decode hits a CUDA illegal memory access when a
+        # request's total length reaches max_model_len (the drafter looks ahead
+        # past the boundary). Clamp per-request max_tokens so speculative
+        # requests stop short of the boundary by the drafter lookahead.
+        spec_cfg = self.cfg.get("vllm_kwargs", {}).get("speculative_config") or {}
+        spec_lookahead = int(spec_cfg.get("num_speculative_tokens", 0))
+        if spec_lookahead > 0:
+            max_model_len = self.cfg["vllm_cfg"]["max_model_len"]
+            base_max_tokens = sampling_params.max_tokens
+            sampling_params = [
+                self._build_sampling_params(
+                    greedy=greedy,
+                    stop_strings=stop_strings,
+                    max_new_tokens=self._spec_decode_max_tokens(
+                        base_max_tokens, int(input_len), max_model_len, spec_lookahead
+                    ),
+                )
+                for input_len in data["input_lengths"].tolist()
+            ]
+
         # verify inputs have correct padding
         verify_right_padding(data, pad_value=self.cfg["_pad_token_id"])
 

nemo_rl/models/megatron/draft/utils.py

@@ -1203,6 +1203,28 @@ def copy_policy_lm_head_to_draft(
         )
 
 
+DRAFT_GRAD_NORM_GROUP = "draft"
+
+
+def register_draft_grad_norm_group() -> None:
+    """Register the 'draft' grad-norm group with Megatron's optimizer.
+
+    Megatron clips parameters in a registered group separately from the main
+    gradient norm (see MegatronOptimizer.clip_grad_norm and the 'mtp'
+    precedent in multi_token_prediction.py), so the draft head's large
+    early-training gradients do not shrink the policy update through the
+    shared global clip. Only called when a draft model is built, so baseline
+    (no-draft) runs keep Megatron's stock clipping behavior.
+    """
+    from megatron.core.optimizer import optimizer as mcore_optimizer
+
+    if DRAFT_GRAD_NORM_GROUP not in mcore_optimizer.SEPARATE_GRAD_NORM_GROUPS:
+        mcore_optimizer.SEPARATE_GRAD_NORM_GROUPS = (
+            *mcore_optimizer.SEPARATE_GRAD_NORM_GROUPS,
+            DRAFT_GRAD_NORM_GROUP,
+        )
+
+
 def build_draft_model(
     model_provider,
     draft_config: dict[str, Any],
@@ -1347,4 +1369,11 @@ def build_draft_model(
         )
         print("[draft] Initialized draft LM head from the policy output layer.")
 
+    # Tag draft params before optimizer construction so
+    # copy_optimizer_param_metadata propagates the group to the distributed
+    # optimizer's shard/fp32 main params and they are clipped separately.
+    register_draft_grad_norm_group()
+    for param in draft_model.parameters():
+        param.grad_norm_group = DRAFT_GRAD_NORM_GROUP
+
     return draft_model

nemo_rl/models/policy/lm_policy.py

@@ -791,6 +791,8 @@ def train(
             aggregated_results["moe_metrics"] = results[0]["moe_metrics"]
         if "mtp_metrics" in results[0]:
             aggregated_results["mtp_metrics"] = results[0]["mtp_metrics"]
+        if "draft_grad_norm" in results[0]:
+            aggregated_results["draft_grad_norm"] = results[0]["draft_grad_norm"]
 
         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

@@ -748,9 +748,15 @@ def train(
                     # (MTP params are tagged only when mtp_detach_heads=True, on the last
                     # pipeline stage). grad_norms_by_group always exists after step().
                     mtp_grad_norm = self.optimizer.grad_norms_by_group.get("mtp")
+                    # Draft params are tagged with their own grad-norm group
+                    # (see build_draft_model) and clipped separately from the
+                    # policy so their large early gradients don't shrink the
+                    # policy update. None when no draft model is attached.
+                    draft_grad_norm = self.optimizer.grad_norms_by_group.get("draft")
                 else:
                     update_successful, grad_norm, num_zeros_in_grad = (True, 0.0, 0.0)
                     mtp_grad_norm = None
+                    draft_grad_norm = None
 
                 pg_collection = get_pg_collection(self.model)
 
@@ -772,6 +778,11 @@ def train(
                 mtp_grad_norm = reduce_max_stat_across_model_parallel_group(
                     mtp_grad_norm, mp_group=pg_collection.mp
                 )
+                # Same for the draft grad norm: the draft model lives on a single
+                # PP stage, so other ranks see None until reduced.
+                draft_grad_norm = reduce_max_stat_across_model_parallel_group(
+                    draft_grad_norm, mp_group=pg_collection.mp
+                )
                 if (
                     not eval_mode
                     and self._first_train_step_forward_pre_hook_disabled
@@ -875,6 +886,8 @@ def train(
         # 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, total_num_microbatches, mtp_grad_norm)
+        if draft_grad_norm is not None:
+            metrics["draft_grad_norm"] = torch.tensor([draft_grad_norm])
 
         # Skip FLOPs estimation when sequence packing is enabled: gbs counts original
         # samples but each packed sequence spans max_total_sequence_length tokens,

tests/unit/models/generation/test_vllm_spec_decode_clamp.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2025, 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from nemo_rl.models.generation.vllm.vllm_worker import BaseVllmGenerationWorker
+
+
+@pytest.mark.parametrize(
+    ("base_max_tokens", "input_len", "max_model_len", "spec_lookahead", "expected"),
+    [
+        (256, 100, 1024, 5, 256),  # clamp inactive: base wins
+        (256, 900, 1024, 5, 118),  # clamp active: 1024 - 900 - 6
+        (256, 1018, 1024, 5, 1),  # at boundary: floor at 1
+        (256, 1050, 1024, 5, 1),  # past boundary: floor at 1
+        (8, 100, 1024, 5, 8),  # base < headroom: base wins
+    ],
+)
+def test_spec_decode_max_tokens_clamp(
+    base_max_tokens, input_len, max_model_len, spec_lookahead, expected
+):
+    assert (
+        BaseVllmGenerationWorker._spec_decode_max_tokens(
+            base_max_tokens, input_len, max_model_len, spec_lookahead
+        )
+        == expected
+    )

tests/unit/models/megatron/test_draft_grad_norm_group.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2025, 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+
+@pytest.mark.mcore
+def test_register_draft_grad_norm_group_is_idempotent_and_preserves_existing():
+    from megatron.core.optimizer import optimizer as mcore_opt
+
+    from nemo_rl.models.megatron.draft.utils import (
+        DRAFT_GRAD_NORM_GROUP,
+        register_draft_grad_norm_group,
+    )
+
+    original = mcore_opt.SEPARATE_GRAD_NORM_GROUPS
+    try:
+        register_draft_grad_norm_group()
+        after_first = mcore_opt.SEPARATE_GRAD_NORM_GROUPS
+        assert DRAFT_GRAD_NORM_GROUP in after_first
+        assert "mtp" in after_first  # not overwritten
+        register_draft_grad_norm_group()
+        assert mcore_opt.SEPARATE_GRAD_NORM_GROUPS == after_first  # no-op
+    finally:
+        mcore_opt.SEPARATE_GRAD_NORM_GROUPS = original