fix(eagle3): isolate draft grad clipping - #3323
Merged
Merged
Conversation
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 <noreply@anthropic.com> Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
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 <noreply@anthropic.com> Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
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) <noreply@anthropic.com> Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
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) <noreply@anthropic.com>
Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
This reverts commit 4dd259f. The decoupled draft learning-rate feature is not needed for now. It is kept in git history (commit 4dd259f) 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) <noreply@anthropic.com> Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
yuekaizhang
changed the title
Contributor
Author
isomap
reviewed
Jul 28, 2026
…yLossFn 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) <noreply@anthropic.com> Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
yuekaizhang
added
the
CI:Lfast
Contributor
Author
|
/ok to test b9dfa51 |
Auto-formatted 3 files that ruff-format reformatted in CI. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Contributor
Author
|
/ok to test a8317ca |
yuki-97
reviewed
Jul 30, 2026
yuki-97
left a comment
yuki-97
left a comment
Contributor
There was a problem hiding this comment.
thanks for the fix, LGTM and let's add some tests to guard.
…m group - 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) <noreply@anthropic.com> Signed-off-by: Yuekai Zhang <zhangyuekai@foxmail.com>
Contributor
Author
|
/ok to test bbdb04c |
yuki-97
approved these changes
Jul 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two fixes for eagle3 draft co-training.
1. Spec decode crashes at the max_model_len boundary
The drafter looks
num_speculative_tokensahead, so when a request reachesmax_model_lenit reads/writes KV slots past the engine buffers → CUDA illegalmemory access. Fix: clamp per-request
max_tokensto leavenum_speculative_tokens + 1headroom, so generation stops before the draftercrosses the boundary. No-op when speculative decoding is off.
2. Draft gradients shrink the policy update
With draft co-training, the freshly-initialized draft head dominates the shared
gradient norm (2–5× the policy-only norm), so the global
max_grad_normcliprescales the policy gradient much smaller than the no-draft baseline — making
the curves incomparable. Fix: clip the draft head in its own grad-norm group so
policy and draft are clipped independently. No-draft runs are unchanged;
train/grad_normis now policy-only andtrain/draft_grad_normis logged.Green Curve: baseline without eagle3
Red Curve: ealge3 + RL with the PR's fix
Yello Curve: eagle3 + RL without the PR's fix