From 47900158df952922fcac68bfd8e3b54c93eb73fd Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Wed, 22 Jul 2026 02:13:03 -0700 Subject: [PATCH 1/8] fix: keep eagle3 spec decode clear of the max_model_len boundary vLLM 0.20 eagle3 speculative decoding hits a CUDA illegal memory access when a request's total length reaches max_model_len: the drafter looks ahead num_speculative_tokens past the current position and reads/writes KV slots beyond the engine buffers. Crash dumps consistently showed requests within num_spec_tokens+1 of the cap (e.g. 8188+1+3 == 8192). Clamp per-request max_tokens to max_model_len - prompt_len - (num_speculative_tokens + 1) so generation length-stops before the drafter can cross the boundary. No-op when speculative decoding is disabled. Give the recipe's vllm max_model_len the same headroom (e.g. 4096 + 4) to keep generation totals identical to the non-speculative baseline. Co-Authored-By: Claude Fable 5 Signed-off-by: Yuekai Zhang --- nemo_rl/models/generation/vllm/vllm_worker.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index a2a59292ca5..0230ab76422 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -687,6 +687,30 @@ 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=max( + 1, + min( + base_max_tokens, + max_model_len - int(input_len) - (spec_lookahead + 1), + ), + ), + ) + for input_len in data["input_lengths"].tolist() + ] + # verify inputs have correct padding verify_right_padding(data, pad_value=self.cfg["_pad_token_id"]) From 3286df3db852c4df13642e5fe4a8b03ff9210d02 Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Wed, 22 Jul 2026 02:13:17 -0700 Subject: [PATCH 2/8] fix: clip draft-model grads in a separate grad-norm group With draft co-training (policy.draft.enabled), the draft head's loss is added to the policy loss and shares one optimizer step. The freshly initialized draft head dominates the global grad norm (measured 2-5x the policy-only norm), so the shared clip at max_grad_norm rescaled the policy gradient 2-5x smaller than an identical no-draft baseline, making eagle3-vs-baseline training curves incomparable. Tag draft params with Megatron's separate grad-norm group mechanism (the 'mtp' precedent): the policy is clipped on its own norm, the draft head on its own, each to max_grad_norm. Tagging happens only when a draft model is built, so no-draft runs keep Megatron's stock clipping byte-for-byte. train/grad_norm is now policy-only; a new train/draft_grad_norm metric reports the draft group. Also detach teacher logits in DraftCrossEntropyLossFn's non-TP fallback to match DistributedCrossEntropy semantics (backward flows only through student logits, never into the policy). Co-Authored-By: Claude Fable 5 Signed-off-by: Yuekai Zhang --- nemo_rl/algorithms/grpo.py | 4 +++ nemo_rl/algorithms/loss/loss_functions.py | 6 +++- nemo_rl/models/megatron/draft/utils.py | 29 +++++++++++++++++++ .../policy/workers/megatron_policy_worker.py | 15 ++++++++++ 4 files changed, 53 insertions(+), 1 deletion(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index a3a725d0cb6..e7801fc140b 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -3002,6 +3002,8 @@ 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() @@ -4377,6 +4379,8 @@ 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(): diff --git a/nemo_rl/algorithms/loss/loss_functions.py b/nemo_rl/algorithms/loss/loss_functions.py index 5b8b0bf9cd0..b94226c229f 100755 --- a/nemo_rl/algorithms/loss/loss_functions.py +++ b/nemo_rl/algorithms/loss/loss_functions.py @@ -92,7 +92,11 @@ def __call__( False, ) else: - teacher_probs = torch.nn.functional.softmax(teacher_logits, dim=-1) + # Match DistributedCrossEntropy semantics: backward propagates only + # through student logits, never into the (policy) teacher. + teacher_probs = torch.nn.functional.softmax( + teacher_logits.detach(), 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) diff --git a/nemo_rl/models/megatron/draft/utils.py b/nemo_rl/models/megatron/draft/utils.py index 19a3912dde6..d5e148f22f8 100644 --- a/nemo_rl/models/megatron/draft/utils.py +++ b/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 diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index cba78d3acd4..98b6832e99f 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -733,9 +733,17 @@ 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) @@ -757,6 +765,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 @@ -860,6 +873,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, From 65078a1cff288dbf126a43eca07cef162e2cc24d Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Wed, 22 Jul 2026 23:42:11 -0700 Subject: [PATCH 3/8] fix: forward draft_grad_norm from policy worker to logged metrics The Megatron policy worker computes the draft model's separate grad norm and sets metrics["draft_grad_norm"] (megatron_policy_worker.py), but Policy.train() only hand-copied a fixed set of top-level worker keys (loss, grad_norm, moe_metrics, mtp_metrics, flops) into aggregated_results. draft_grad_norm was silently dropped, so it never reached wandb/tensorboard even though the draft grads are clipped in their own group. Forward draft_grad_norm the same way mtp_metrics is forwarded, so train/draft_grad_norm is logged for any run with policy.draft.enabled=true. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Yuekai Zhang --- nemo_rl/models/policy/lm_policy.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/nemo_rl/models/policy/lm_policy.py b/nemo_rl/models/policy/lm_policy.py index 397b4e086b5..0c858029a29 100644 --- a/nemo_rl/models/policy/lm_policy.py +++ b/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 From 4dd259fa5eda355bc4daf8774a5c75e9086839f0 Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Thu, 23 Jul 2026 02:00:58 -0700 Subject: [PATCH 4/8] feat(megatron): independent LR + train/draft_lr logging for eagle3 draft Add an optional decoupled learning rate for the Eagle draft submodule so a from-scratch draft can train at a larger LR than the fine-tuned policy. - policy.draft.lr / min_lr (DraftConfig): when lr is set, draft params get a dedicated optimizer param group at that peak LR; unset keeps the policy LR. - setup.py: _DraftLROverrideProvider passed to setup_optimizer builds a ParamGroupOverride({"max_lr": draft_lr}) matching draft params by the grad_norm_group=="draft" attribute (robust to Float16Module/DDP name prefixes present at optimizer-build time). The scheduler drives each group by its own max_lr, same path as decoupled/embedding LR. - megatron_policy_worker.py: log the draft group's scheduled LR as train/draft_lr alongside train/lr. - grpo.py: add draft_lr to the per-microbatch mean-allowlist (otherwise unlisted keys are summed, inflating the logged value). - add recipe grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr.yaml. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Yuekai Zhang --- ...megatron-eagle3-scratch-draft-draftlr.yaml | 64 +++++++++++++++++++ nemo_rl/algorithms/grpo.py | 2 + nemo_rl/models/megatron/setup.py | 58 +++++++++++++++++ nemo_rl/models/policy/__init__.py | 7 ++ .../policy/workers/megatron_policy_worker.py | 26 ++++++++ 5 files changed, 157 insertions(+) create mode 100644 examples/configs/recipes/llm/grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr.yaml diff --git a/examples/configs/recipes/llm/grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr.yaml b/examples/configs/recipes/llm/grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr.yaml new file mode 100644 index 00000000000..4fc3661e84f --- /dev/null +++ b/examples/configs/recipes/llm/grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr.yaml @@ -0,0 +1,64 @@ +defaults: ../../grpo_math_1B.yaml + +# Identical to grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft.yaml, except +# the from-scratch draft trains with its OWN peak learning rate (policy.draft.lr). +# The draft params get a dedicated optimizer param group whose max_lr comes from +# policy.draft.lr, following the same warmup/decay schedule as the policy; every +# other param keeps the policy LR (megatron_cfg.optimizer.lr). This lets the +# random-init draft learn faster than the fine-tuned 8B policy. +# The draft LR schedule is logged live as train/draft_lr, and the peak is encoded +# in the wandb run name (draftlr1e-4) for tracking. + +checkpointing: + checkpoint_dir: results/grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr1e-4 +policy: + model_name: Qwen/Qwen3-8B-Base + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 4096 + make_sequence_length_divisible_by: 4 + dtensor_cfg: + enabled: false + sequence_packing: + enabled: false + megatron_cfg: + enabled: true + tensor_model_parallel_size: 4 + sequence_parallel: true + apply_rope_fusion: false + activation_checkpointing: true + defer_fp32_logits: true + draft: + enabled: true + # Random-init draft: architecture derived from the policy model + # (1 decoder layer, full 151936 vocab, aux layers auto-selected). + model_name: null + loss_weight: 1.0 + num_layers: 1 + aux_layer_indices: null + # Decoupled draft learning rate (from-scratch draft needs a larger LR than + # the fine-tuned policy, whose megatron_cfg.optimizer.lr is 5e-6). Tunable. + lr: 1.0e-4 + generation: + vllm_cfg: + tensor_parallel_size: 4 + gpu_memory_utilization: 0.6 + max_model_len: 4100 # 4096 + num_speculative_tokens + 1 (drafter headroom) + vllm_kwargs: + speculative_config: + method: "eagle3" + # Local config dir (no weights): same architecture as + # AngelSlim/Qwen3-8B_eagle3 but draft_vocab_size == vocab_size so shapes + # match the scratch megatron draft. vLLM dummy-loads it and the first + # refit pushes the random-init weights from the trainer. + model: /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yuekaiz/speculative_rl/RL/results/eagle3-scratch-draft-8b-vllm-config + num_speculative_tokens: 3 + draft_tensor_parallel_size: 1 +logger: + wandb_enabled: true + tensorboard_enabled: true + wandb: + project: nemo-rl + name: grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr1e-4 +cluster: + gpus_per_node: 8 diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index e7801fc140b..152029d4fc8 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -3024,6 +3024,7 @@ def grpo_train( ) elif k in { "lr", + "draft_lr", "wd", "reward", "filtered_reward", @@ -4396,6 +4397,7 @@ def async_grpo_train( ) elif k in { "lr", + "draft_lr", "wd", "reward", "global_valid_seqs", diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index a88cd334114..857ef9128e6 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -18,6 +18,7 @@ import threading import time import warnings +from dataclasses import dataclass from typing import Any, Callable, Optional, TypeVar import torch @@ -38,10 +39,18 @@ DistributedInitConfig, LoggerConfig, OptimizerConfig, + OptimizerConfigOverrideProvider, + OptimizerConfigOverrideProviderContext, SchedulerConfig, TokenizerConfig, TrainingConfig, ) +from megatron.core.optimizer import ( + ParamGroupOverride, + ParamKey, + ParamPredicate, + get_standard_config_overrides, +) from megatron.bridge.training.initialize import ( initialize_megatron, set_jit_fusion_options, @@ -68,6 +77,45 @@ _HF_CONFIG_PATCHED = False +def _is_draft_param(param: "torch.nn.Parameter") -> bool: + """True for Eagle draft params, tagged with grad_norm_group == 'draft'. + + Matching by this attribute (set in build_draft_model) rather than by name is + robust to the module wrapper prefixes (Float16Module/DDP) present on the + param names at optimizer-construction time. + """ + return getattr(param, "grad_norm_group", None) == DRAFT_GRAD_NORM_GROUP + + +@dataclass +class _DraftLROverrideProvider(OptimizerConfigOverrideProvider): + """Give the Eagle draft submodule its own peak learning rate. + + Draft params are placed in a dedicated optimizer param group whose + ``max_lr``/``min_lr`` come from ``policy.draft.lr``/``policy.draft.min_lr``; + all other params keep the policy LR. The scheduler drives each group by its + own ``max_lr``, so the draft follows the same warmup/decay shape scaled to + its own peak — the same mechanism Megatron uses for decoupled/embedding LRs. + This lets a from-scratch draft train at a larger LR than the fine-tuned + policy. Standard weight-decay overrides are preserved. + """ + + draft_lr: float + draft_min_lr: float | None = None + + def build_config_overrides( + self, context: OptimizerConfigOverrideProviderContext + ) -> dict[ParamKey, ParamGroupOverride] | None: + overrides = get_standard_config_overrides(config=context.optimizer_config) or {} + draft_override: ParamGroupOverride = {"max_lr": self.draft_lr} + if self.draft_min_lr is not None: + draft_override["min_lr"] = self.draft_min_lr + draft_key = ParamKey( + predicate=ParamPredicate(name=DRAFT_GRAD_NORM_GROUP, fn=_is_draft_param) + ) + overrides[draft_key] = draft_override + return overrides + def _patch_hf_config_double_instantiation(): """Patch HF config classes whose __post_init__ fails with Megatron's recursive instantiation. @@ -134,6 +182,7 @@ def _safe_post_init(self, **kwargs): from nemo_rl.models.megatron.community_import import import_model_from_hf_name from nemo_rl.models.megatron.config import ModelAndOptimizerState, RuntimeConfig from nemo_rl.models.megatron.draft.utils import ( + DRAFT_GRAD_NORM_GROUP, build_draft_model, find_draft_owner_chunk, get_attached_draft_model, @@ -1377,11 +1426,20 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: ) if load_optimizer: + # Give the draft submodule its own peak LR when policy.draft.lr is set. + # None keeps the default provider (draft shares the policy LR). + optimizer_config_override_provider = None + if draft_enabled and policy_cfg["draft"].get("lr") is not None: + optimizer_config_override_provider = _DraftLROverrideProvider( + draft_lr=policy_cfg["draft"]["lr"], + draft_min_lr=policy_cfg["draft"].get("min_lr"), + ) optimizer, scheduler = setup_optimizer( optimizer_config=megatron_cfg.optimizer, scheduler_config=megatron_cfg.scheduler, model=model, use_gloo_process_groups=megatron_cfg.dist.use_gloo_process_groups, + optimizer_config_override_provider=optimizer_config_override_provider, ) else: optimizer = None diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index fcda0190e36..432fcf98ffe 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -457,6 +457,13 @@ class DraftConfig(TypedDict): loss_weight: NotRequired[float] num_layers: NotRequired[int | None] aux_layer_indices: NotRequired[list[int] | None] + # Decoupled learning rate for the draft submodule. When lr is set, the draft + # params train in their own optimizer param group with this peak LR (and + # min_lr, if given), following the same warmup/decay schedule as the policy. + # Useful when the draft is trained from scratch and needs a larger LR than + # the (fine-tuned) policy. When lr is unset/None, the draft shares the policy LR. + lr: NotRequired[float | None] + min_lr: NotRequired[float | None] class TokenizerConfig(TypedDict): diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 98b6832e99f..577460b4b7e 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -140,6 +140,24 @@ def _model_self_packs_for_cp(model: Any) -> bool: return any(isinstance(chunk, Qwen3VLModel) for chunk in chunks) +def _draft_scheduled_lr(optimizer: Any, scheduler: Any) -> Optional[float]: + """Current scheduled LR of the Eagle draft param group, or None if absent. + + Draft params carry ``grad_norm_group == "draft"`` (see build_draft_model), so + they live in their own optimizer param group(s) when ``policy.draft.lr`` is + set. Both draft groups (wd / no-wd) share the same ``max_lr``, so the first + match yields the draft LR. Returns None when there is no dedicated draft + group (draft disabled, or draft.lr unset so the draft shares the policy LR). + """ + for group in optimizer.param_groups: + if any( + getattr(p, "grad_norm_group", None) == "draft" + for p in group.get("params", []) + ): + return scheduler.get_lr(group) + return None + + # Classes with @ray.remote can't be inherited from, so we split the implementation out. # This is useful when using worker extension classes. class MegatronPolicyWorkerImpl( @@ -804,6 +822,11 @@ def train( curr_wd = self.scheduler.get_wd() loss_metrics["lr"] = curr_lr loss_metrics["wd"] = curr_wd + draft_lr = _draft_scheduled_lr( + self.optimizer, self.scheduler + ) + if draft_lr is not None: + loss_metrics["draft_lr"] = draft_lr loss_metrics["global_valid_seqs"] = global_valid_seqs.item() loss_metrics["global_valid_toks"] = global_valid_toks.item() mb_losses.append(loss_metrics["loss"]) @@ -1313,6 +1336,7 @@ def _finish_train_step_body(self, state: dict[str, Any]) -> dict[str, Any]: # the value of THIS step, not the next one. (terrykong, #2683:832). curr_lr = self.scheduler.get_lr(self.optimizer.param_groups[0]) curr_wd = self.scheduler.get_wd() + draft_lr = _draft_scheduled_lr(self.optimizer, self.scheduler) # Scheduler increment matches sync path's ``increment=gbs``. self.scheduler.step(increment=state["gbs"]) @@ -1375,6 +1399,8 @@ def _scale_metric(name: str, value: Any) -> Any: out[k] = _scale_metric(k, v) out["lr"] = curr_lr out["wd"] = curr_wd + if draft_lr is not None: + out["draft_lr"] = draft_lr out["global_valid_seqs"] = global_valid_seqs_f out["global_valid_toks"] = global_valid_toks_f rescaled_metrics.append(out) From 03f9e04d013e87e3671ced19e719c773f81ea89f Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Thu, 23 Jul 2026 02:01:19 -0700 Subject: [PATCH 5/8] revert: independent LR + train/draft_lr logging for eagle3 draft This reverts commit 4dd259fa5. The decoupled draft learning-rate feature is not needed for now. It is kept in git history (commit 4dd259fa5) for future reference and can be cherry-picked back if we revisit per-submodule LRs for from-scratch draft training. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Yuekai Zhang --- ...megatron-eagle3-scratch-draft-draftlr.yaml | 64 ------------------- nemo_rl/algorithms/grpo.py | 2 - nemo_rl/models/megatron/setup.py | 58 ----------------- nemo_rl/models/policy/__init__.py | 7 -- .../policy/workers/megatron_policy_worker.py | 26 -------- 5 files changed, 157 deletions(-) delete mode 100644 examples/configs/recipes/llm/grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr.yaml diff --git a/examples/configs/recipes/llm/grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr.yaml b/examples/configs/recipes/llm/grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr.yaml deleted file mode 100644 index 4fc3661e84f..00000000000 --- a/examples/configs/recipes/llm/grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr.yaml +++ /dev/null @@ -1,64 +0,0 @@ -defaults: ../../grpo_math_1B.yaml - -# Identical to grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft.yaml, except -# the from-scratch draft trains with its OWN peak learning rate (policy.draft.lr). -# The draft params get a dedicated optimizer param group whose max_lr comes from -# policy.draft.lr, following the same warmup/decay schedule as the policy; every -# other param keeps the policy LR (megatron_cfg.optimizer.lr). This lets the -# random-init draft learn faster than the fine-tuned 8B policy. -# The draft LR schedule is logged live as train/draft_lr, and the peak is encoded -# in the wandb run name (draftlr1e-4) for tracking. - -checkpointing: - checkpoint_dir: results/grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr1e-4 -policy: - model_name: Qwen/Qwen3-8B-Base - train_micro_batch_size: 1 - logprob_batch_size: 1 - max_total_sequence_length: 4096 - make_sequence_length_divisible_by: 4 - dtensor_cfg: - enabled: false - sequence_packing: - enabled: false - megatron_cfg: - enabled: true - tensor_model_parallel_size: 4 - sequence_parallel: true - apply_rope_fusion: false - activation_checkpointing: true - defer_fp32_logits: true - draft: - enabled: true - # Random-init draft: architecture derived from the policy model - # (1 decoder layer, full 151936 vocab, aux layers auto-selected). - model_name: null - loss_weight: 1.0 - num_layers: 1 - aux_layer_indices: null - # Decoupled draft learning rate (from-scratch draft needs a larger LR than - # the fine-tuned policy, whose megatron_cfg.optimizer.lr is 5e-6). Tunable. - lr: 1.0e-4 - generation: - vllm_cfg: - tensor_parallel_size: 4 - gpu_memory_utilization: 0.6 - max_model_len: 4100 # 4096 + num_speculative_tokens + 1 (drafter headroom) - vllm_kwargs: - speculative_config: - method: "eagle3" - # Local config dir (no weights): same architecture as - # AngelSlim/Qwen3-8B_eagle3 but draft_vocab_size == vocab_size so shapes - # match the scratch megatron draft. vLLM dummy-loads it and the first - # refit pushes the random-init weights from the trainer. - model: /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_nemorl/users/yuekaiz/speculative_rl/RL/results/eagle3-scratch-draft-8b-vllm-config - num_speculative_tokens: 3 - draft_tensor_parallel_size: 1 -logger: - wandb_enabled: true - tensorboard_enabled: true - wandb: - project: nemo-rl - name: grpo-qwen3-8b-base-1n8g-megatron-eagle3-scratch-draft-draftlr1e-4 -cluster: - gpus_per_node: 8 diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 152029d4fc8..e7801fc140b 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -3024,7 +3024,6 @@ def grpo_train( ) elif k in { "lr", - "draft_lr", "wd", "reward", "filtered_reward", @@ -4397,7 +4396,6 @@ def async_grpo_train( ) elif k in { "lr", - "draft_lr", "wd", "reward", "global_valid_seqs", diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 857ef9128e6..a88cd334114 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -18,7 +18,6 @@ import threading import time import warnings -from dataclasses import dataclass from typing import Any, Callable, Optional, TypeVar import torch @@ -39,18 +38,10 @@ DistributedInitConfig, LoggerConfig, OptimizerConfig, - OptimizerConfigOverrideProvider, - OptimizerConfigOverrideProviderContext, SchedulerConfig, TokenizerConfig, TrainingConfig, ) -from megatron.core.optimizer import ( - ParamGroupOverride, - ParamKey, - ParamPredicate, - get_standard_config_overrides, -) from megatron.bridge.training.initialize import ( initialize_megatron, set_jit_fusion_options, @@ -77,45 +68,6 @@ _HF_CONFIG_PATCHED = False -def _is_draft_param(param: "torch.nn.Parameter") -> bool: - """True for Eagle draft params, tagged with grad_norm_group == 'draft'. - - Matching by this attribute (set in build_draft_model) rather than by name is - robust to the module wrapper prefixes (Float16Module/DDP) present on the - param names at optimizer-construction time. - """ - return getattr(param, "grad_norm_group", None) == DRAFT_GRAD_NORM_GROUP - - -@dataclass -class _DraftLROverrideProvider(OptimizerConfigOverrideProvider): - """Give the Eagle draft submodule its own peak learning rate. - - Draft params are placed in a dedicated optimizer param group whose - ``max_lr``/``min_lr`` come from ``policy.draft.lr``/``policy.draft.min_lr``; - all other params keep the policy LR. The scheduler drives each group by its - own ``max_lr``, so the draft follows the same warmup/decay shape scaled to - its own peak — the same mechanism Megatron uses for decoupled/embedding LRs. - This lets a from-scratch draft train at a larger LR than the fine-tuned - policy. Standard weight-decay overrides are preserved. - """ - - draft_lr: float - draft_min_lr: float | None = None - - def build_config_overrides( - self, context: OptimizerConfigOverrideProviderContext - ) -> dict[ParamKey, ParamGroupOverride] | None: - overrides = get_standard_config_overrides(config=context.optimizer_config) or {} - draft_override: ParamGroupOverride = {"max_lr": self.draft_lr} - if self.draft_min_lr is not None: - draft_override["min_lr"] = self.draft_min_lr - draft_key = ParamKey( - predicate=ParamPredicate(name=DRAFT_GRAD_NORM_GROUP, fn=_is_draft_param) - ) - overrides[draft_key] = draft_override - return overrides - def _patch_hf_config_double_instantiation(): """Patch HF config classes whose __post_init__ fails with Megatron's recursive instantiation. @@ -182,7 +134,6 @@ def _safe_post_init(self, **kwargs): from nemo_rl.models.megatron.community_import import import_model_from_hf_name from nemo_rl.models.megatron.config import ModelAndOptimizerState, RuntimeConfig from nemo_rl.models.megatron.draft.utils import ( - DRAFT_GRAD_NORM_GROUP, build_draft_model, find_draft_owner_chunk, get_attached_draft_model, @@ -1426,20 +1377,11 @@ def composed_peft_hook(model: list[MegatronModule]) -> list[MegatronModule]: ) if load_optimizer: - # Give the draft submodule its own peak LR when policy.draft.lr is set. - # None keeps the default provider (draft shares the policy LR). - optimizer_config_override_provider = None - if draft_enabled and policy_cfg["draft"].get("lr") is not None: - optimizer_config_override_provider = _DraftLROverrideProvider( - draft_lr=policy_cfg["draft"]["lr"], - draft_min_lr=policy_cfg["draft"].get("min_lr"), - ) optimizer, scheduler = setup_optimizer( optimizer_config=megatron_cfg.optimizer, scheduler_config=megatron_cfg.scheduler, model=model, use_gloo_process_groups=megatron_cfg.dist.use_gloo_process_groups, - optimizer_config_override_provider=optimizer_config_override_provider, ) else: optimizer = None diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 432fcf98ffe..fcda0190e36 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/nemo_rl/models/policy/__init__.py @@ -457,13 +457,6 @@ class DraftConfig(TypedDict): loss_weight: NotRequired[float] num_layers: NotRequired[int | None] aux_layer_indices: NotRequired[list[int] | None] - # Decoupled learning rate for the draft submodule. When lr is set, the draft - # params train in their own optimizer param group with this peak LR (and - # min_lr, if given), following the same warmup/decay schedule as the policy. - # Useful when the draft is trained from scratch and needs a larger LR than - # the (fine-tuned) policy. When lr is unset/None, the draft shares the policy LR. - lr: NotRequired[float | None] - min_lr: NotRequired[float | None] class TokenizerConfig(TypedDict): diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 577460b4b7e..98b6832e99f 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -140,24 +140,6 @@ def _model_self_packs_for_cp(model: Any) -> bool: return any(isinstance(chunk, Qwen3VLModel) for chunk in chunks) -def _draft_scheduled_lr(optimizer: Any, scheduler: Any) -> Optional[float]: - """Current scheduled LR of the Eagle draft param group, or None if absent. - - Draft params carry ``grad_norm_group == "draft"`` (see build_draft_model), so - they live in their own optimizer param group(s) when ``policy.draft.lr`` is - set. Both draft groups (wd / no-wd) share the same ``max_lr``, so the first - match yields the draft LR. Returns None when there is no dedicated draft - group (draft disabled, or draft.lr unset so the draft shares the policy LR). - """ - for group in optimizer.param_groups: - if any( - getattr(p, "grad_norm_group", None) == "draft" - for p in group.get("params", []) - ): - return scheduler.get_lr(group) - return None - - # Classes with @ray.remote can't be inherited from, so we split the implementation out. # This is useful when using worker extension classes. class MegatronPolicyWorkerImpl( @@ -822,11 +804,6 @@ def train( curr_wd = self.scheduler.get_wd() loss_metrics["lr"] = curr_lr loss_metrics["wd"] = curr_wd - draft_lr = _draft_scheduled_lr( - self.optimizer, self.scheduler - ) - if draft_lr is not None: - loss_metrics["draft_lr"] = draft_lr loss_metrics["global_valid_seqs"] = global_valid_seqs.item() loss_metrics["global_valid_toks"] = global_valid_toks.item() mb_losses.append(loss_metrics["loss"]) @@ -1336,7 +1313,6 @@ def _finish_train_step_body(self, state: dict[str, Any]) -> dict[str, Any]: # the value of THIS step, not the next one. (terrykong, #2683:832). curr_lr = self.scheduler.get_lr(self.optimizer.param_groups[0]) curr_wd = self.scheduler.get_wd() - draft_lr = _draft_scheduled_lr(self.optimizer, self.scheduler) # Scheduler increment matches sync path's ``increment=gbs``. self.scheduler.step(increment=state["gbs"]) @@ -1399,8 +1375,6 @@ def _scale_metric(name: str, value: Any) -> Any: out[k] = _scale_metric(k, v) out["lr"] = curr_lr out["wd"] = curr_wd - if draft_lr is not None: - out["draft_lr"] = draft_lr out["global_valid_seqs"] = global_valid_seqs_f out["global_valid_toks"] = global_valid_toks_f rescaled_metrics.append(out) From 5b82d1f1aa043e67d7aa2df0ca94e3ace3bc1ddd Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Tue, 28 Jul 2026 19:48:38 -0700 Subject: [PATCH 6/8] nit: remove redundant .detach() on teacher_logits in DraftCrossEntropyLossFn teacher_logits is derived from logits.detach() in utils.py before reaching this loss path, so the additional .detach() here is redundant. Keep gradient-boundary ownership in one place. Co-Authored-By: Claude Sonnet 4.6 (1M context) Signed-off-by: Yuekai Zhang --- nemo_rl/algorithms/loss/loss_functions.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nemo_rl/algorithms/loss/loss_functions.py b/nemo_rl/algorithms/loss/loss_functions.py index b94226c229f..7e064fd2ef2 100755 --- a/nemo_rl/algorithms/loss/loss_functions.py +++ b/nemo_rl/algorithms/loss/loss_functions.py @@ -92,10 +92,10 @@ def __call__( False, ) else: - # Match DistributedCrossEntropy semantics: backward propagates only - # through student logits, never into the (policy) teacher. + # teacher_logits is already detached at the call site (utils.py); + # match DistributedCrossEntropy semantics. teacher_probs = torch.nn.functional.softmax( - teacher_logits.detach(), dim=-1 + 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) From a8317ca4682df648a19b3116da4f94ebd71ad17c Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Tue, 28 Jul 2026 23:21:26 -0700 Subject: [PATCH 7/8] style: apply ruff-format to fix CI lint failure Auto-formatted 3 files that ruff-format reformatted in CI. Co-Authored-By: Claude Sonnet 4.6 (1M context) Signed-off-by: Yuekai Zhang --- nemo_rl/algorithms/grpo.py | 8 ++++++-- nemo_rl/algorithms/loss/loss_functions.py | 4 +--- nemo_rl/models/policy/workers/megatron_policy_worker.py | 4 +--- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index a7462038ddd..82c9cb19481 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -3193,7 +3193,9 @@ def grpo_train( {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["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() @@ -4599,7 +4601,9 @@ def async_grpo_train( {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["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(): diff --git a/nemo_rl/algorithms/loss/loss_functions.py b/nemo_rl/algorithms/loss/loss_functions.py index 7e064fd2ef2..3362f238073 100755 --- a/nemo_rl/algorithms/loss/loss_functions.py +++ b/nemo_rl/algorithms/loss/loss_functions.py @@ -94,9 +94,7 @@ def __call__( 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 - ) + 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) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index a2807fbe86a..36fa12ea384 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -752,9 +752,7 @@ def train( # (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" - ) + 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 From bbdb04cc82f9de142bcc85c38da5d3ee5c4b12de Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Wed, 29 Jul 2026 23:33:59 -0700 Subject: [PATCH 8/8] test(eagle3): add unit tests for spec-decode clamp and draft grad-norm group MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Extract _spec_decode_max_tokens as a static method on BaseVllmGenerationWorker so the +1 boundary reservation and max(1,…) floor are testable without a live vLLM engine. - Add parametrized test covering clamp-inactive, clamp-active, at-boundary, past-boundary, and base-wins cases. - Add mcore-marked test for register_draft_grad_norm_group verifying idempotency and that the pre-existing 'mtp' entry is preserved. Co-Authored-By: Claude Sonnet 4.6 (1M context) Signed-off-by: Yuekai Zhang --- nemo_rl/models/generation/vllm/vllm_worker.py | 24 +++++++++--- .../generation/test_vllm_spec_decode_clamp.py | 38 +++++++++++++++++++ .../megatron/test_draft_grad_norm_group.py | 36 ++++++++++++++++++ 3 files changed, 92 insertions(+), 6 deletions(-) create mode 100644 tests/unit/models/generation/test_vllm_spec_decode_clamp.py create mode 100644 tests/unit/models/megatron/test_draft_grad_norm_group.py diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index 4c8522833f0..ff92a42c116 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/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. @@ -820,12 +836,8 @@ def generate( self._build_sampling_params( greedy=greedy, stop_strings=stop_strings, - max_new_tokens=max( - 1, - min( - base_max_tokens, - max_model_len - int(input_len) - (spec_lookahead + 1), - ), + 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() diff --git a/tests/unit/models/generation/test_vllm_spec_decode_clamp.py b/tests/unit/models/generation/test_vllm_spec_decode_clamp.py new file mode 100644 index 00000000000..6baa9b416ac --- /dev/null +++ b/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 + ) diff --git a/tests/unit/models/megatron/test_draft_grad_norm_group.py b/tests/unit/models/megatron/test_draft_grad_norm_group.py new file mode 100644 index 00000000000..094d4d7586e --- /dev/null +++ b/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