feat(speculative): add P-EAGLE parallel-drafting training for EAGLE-3 drafts - #2376
Merged
HuiyingLi merged 18 commits intoJun 4, 2026
Merged
Conversation
… drafts P-EAGLE (https://arxiv.org/abs/2602.01469) trains the EAGLE-3 draft to predict all K tokens in a single forward instead of EAGLE-3's autoregressive TTT unroll: masked multi-token-prediction depths are conditioned on a fixed learnable mask_hidden placeholder and the reserved ptd_token_id token rather than the previous depth's output. Implementation (reuses the EAGLE-3 cache_hidden diagonal-extension attention, which is numerically faithful to vLLM's single-base-position decode): - draft_llama: register a learnable mask_hidden Parameter of shape [1, num_aux_hidden_states * target_hidden_size] under config.parallel_drafting, plus masked_projected_hidden() to project it through fc/fc_norm. - core: PEagleTrainerModule runs the parallel masked-depth loss, sharing the draft-vocab projection, 0.8**d weighted-mean schedule, and metrics with Eagle3TrainerModule. - train_eagle3: parallel_drafting / num_draft_tokens / ptd_token_id recipe args; ptd_token_id is written into the draft config.json. Checkpoint format matches vLLM's parallel-drafting loader (vllm-project/vllm#32887, v0.16.0): the mask_hidden tensor is saved under the key 'mask_hidden' and ptd_token_id rides in config.json, so the head loads into vLLM unchanged. Inference is vLLM-only today; serve_sglang now rejects a parallel_drafting head with an actionable error (SGLang support tracked in sgl-project/sglang#23171). Tests: tests/unit_tests/speculative/test_peagle.py (mask_hidden key/shape, mask inputs vs no-recurrence, K=1 == EAGLE-3 step 0, loss decreases, state-dict round-trip) and a serve_sglang rejection test. CPU-only. Signed-off-by: khazic <khazzz1c@gmail.com>
… num_draft_tokens
recipe_cfg.get('num_draft_tokens', recipe_cfg.ttt_steps) evaluates the
recipe_cfg.ttt_steps default argument unconditionally (Python always computes a
dict.get default), so a P-EAGLE config that correctly omits ttt_steps raised
AttributeError in setup() before any training step. Resolve through a nested
.get with a literal fallback (num_draft_tokens -> ttt_steps -> 4) so the missing
key is never touched. Add a regression test that loads the committed P-EAGLE
example config and asserts it carries no ttt_steps and still resolves K.
Signed-off-by: khazic <khazzz1c@gmail.com>
masked_projected_hidden() and the batch hidden_states are both outputs of the same model.fc projection, so they already share the compute dtype; the .to(hidden_states.dtype) was a no-op. Removing it matches Eagle3TrainerModule, which does not cast. No behavior change (unit tests unchanged). Signed-off-by: khazic <khazzz1c@gmail.com>
khazic
requested review from
HuiyingLi,
ZhiyuLi-Nvidia,
adil-a,
akoumpa,
athitten,
hemildesai,
pthombre and
zyzhou5
as code owners
June 1, 2026 12:53
Contributor
|
/ok to test 76450e0 |
LlamaRotaryEmbedding.forward keyed only on position_ids.shape[-1] (the sequence length) and returned cos_cache[:seq_len], silently ignoring the position values. Any non-contiguous position_ids therefore received the wrong rotary phase: EAGLE TTT depth offsets (arange(seq_len) + step_idx), packed sequences, and context parallelism. In particular EAGLE's intended per-step phase shift was a no-op. Gather cos/sin by the actual position values and size the cache to the highest requested position (which can exceed seq_len for TTT offsets). For the common position_ids == arange(seq_len) case the gather is numerically identical to the previous slice, so standard training/inference is unchanged. The fused TE path, which consumes raw angles by sequence index, keeps the legacy contiguous slice. Add tests/unit_tests/models/llama/test_rope_utils.py covering the arange no-op, the offset case, non-contiguous gather, and out-of-seq_len positions. Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
|
/ok to test 5767091 |
Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
|
/ok to test 8a8a7bb |
Signed-off-by: thyways <2484113689@qq.com>
fix(speculative): add Qwen P-EAGLE training fixes
Contributor
|
/ok to test 788f917 |
Contributor
|
/ok to test 1f91013 |
Contributor
|
/ok to test 1689989 |
…arallel-drafting Signed-off-by: khazic <khazzz1c@gmail.com> # Conflicts: # nemo_automodel/recipes/llm/train_eagle3.py
Contributor
|
/ok to test 5dbeddb |
Fixes the 9 failing L0 CPU unit tests on this branch. 1. draft_llama: the P-EAGLE flex attention was a module-level torch.compile(flex_attention, mode="max-autotune-no-cudagraphs") called unconditionally. Inductor's flex backend is unavailable on CPU and raises InductorError, so all six test_peagle cases failed. Dispatch to the compiled callable only on CUDA and fall back to eager flex_attention on CPU (correct, just slower) -- the compiled callable stays lazy so CPU import is free. 2. eagle3 token mapping: wrapping the dataloader in tqdm(dataloader) triggers a second __iter__ on loaders without __len__, double-scanning the data. This broke the three test_eagle3_token_cache "scan exactly once" assertions and is a real bug for single-pass / streaming loaders. Drive the progress bar manually and iterate the dataloader directly so it is scanned exactly once. Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
|
/ok to test 92e0f80 |
Contributor
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # dst=Path('./cache/dataset/perfectblend-qwen3-8b-regen-messages'); dst.mkdir(parents=True, exist_ok=True); \ | ||
| # files=sorted(f for f in list_repo_files(repo, repo_type='dataset') if f.startswith('data/train-') and f.endswith('.parquet')); \ | ||
| # [pd.read_parquet(hf_hub_download(repo_id=repo, repo_type='dataset', filename=f)).rename(columns={'conversations':'messages'}).to_parquet(dst / Path(f).name, index=False) for f in files]; \ | ||
| # print(f'wrote{len(files)} shards to {dst}')" |
Contributor
There was a problem hiding this comment.
Nit: missing space in the f-string — this would print wrote5 shards instead of wrote 5 shards.
Suggested change
| # print(f'wrote{len(files)} shards to {dst}')" | |
| # print(f'wrote {len(files)} shards to {dst}')" |
| lr: 1e-4 | ||
| betas: [0.9, 0.95] | ||
| weight_decay: 0 | ||
| warmup_ratio: 0.015 No newline at end of file |
Contributor
There was a problem hiding this comment.
Nit: missing newline at end of file.
The P-EAGLE draft forward runs flex_attention, whose autograd is not implemented on CPU in the CI torch build -- even the forward errors with 'FlexAttention does not support backward on CPU' once the inputs require grad. The six tests that drive the draft forward built the model on CPU, so they failed in both the CPU and GPU unit-test jobs. Make the draft / trainer / batch device-aware (CUDA when available) and skip those six tests when CUDA is unavailable; the pure-tensor COD, mask-mod, draft-contract, and config tests stay on CPU. Locally (no CUDA) the suite is 13 passed, 6 skipped. Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
|
/ok to test 927dfbf |
…arallel-drafting Signed-off-by: khazic <khazzz1c@gmail.com> # Conflicts: # nemo_automodel/components/speculative/eagle/__init__.py
Contributor
|
/ok to test 2da9eea |
HuiyingLi
approved these changes
Jun 4, 2026
24 tasks
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
4 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.
What does this PR do ?
Adds P-EAGLE (parallel-drafting EAGLE-3) training to the speculative-decoding
stack. The EAGLE-3 draft learns to predict all
Ktokens in a single forwardpass (parallel multi-token prediction) instead of EAGLE-3's autoregressive
test-time-training (TTT) unroll. Checkpoints are written in vLLM's
parallel-drafting layout so they load into vLLM ≥ 0.16 unchanged.
P-EAGLE paper: https://arxiv.org/abs/2602.01469 · vLLM runtime:
vllm-project/vllm#32887 (released in v0.16.0).
Changelog
components/speculative/eagle/draft_llama.py: register a learnablemask_hiddenplaceholder of shape[1, num_aux_hidden_states * target_hidden_size]under
config.parallel_drafting(guarded, same pattern as the EAGLE-3.1fc_norm/norm_outputtoggles); addmasked_projected_hidden().components/speculative/eagle/core.py: newPEagleTrainerModule. Depth 0(NTP) consumes the real token + projected target aux hidden states; depths
≥ 1 (MTP) consume the masked
ptd_token_idtoken and the fixed learnablemask_hiddenplaceholder, with no recurrence. Reuses the EAGLE-3cache_hiddenattention, the draft-vocab projection, the0.8**dlossschedule, and the metrics.
recipes/llm/train_eagle3.py:parallel_drafting/num_draft_tokens/ptd_token_idrecipe args.ptd_token_idis required whenparallel_drafting=True(no silent default) and range-checked; it is writteninto the draft
config.json.components/speculative/serve_sglang.py: reject a P-EAGLE head with anactionable error — SGLang cannot serve it yet (tracked in
[Feature] Support P-EAGLE (Parallel EAGLE) speculative decoding sgl-project/sglang#23171).
examples/speculative/p-eagle/llama_peagle_mvp.yaml: example config (uses areserved Llama-3 special token for
ptd_token_id).tests/unit_tests/speculative/test_peagle.py(15) + a serve_sglangrejection test.
Checkpoint contract (vLLM-loadable)
vLLM's parallel-drafting loader keys on the substring
mask_hiddenand readsptd_token_idfrom the draft config. This PR saves themask_hiddentensorunder that key and writes
ptd_token_idintoconfig.json; everything else isa standard EAGLE-3 draft checkpoint. Inference is vLLM-only today (SGLang
support is upstream-tracked, not merged), documented in the example YAML and
enforced by the serve guard.
Equivalence to vLLM inference (parity)
test_peagle_parallel_forward_matches_expanded_referencebuilds an independentflat expanded parallel forward — the
[context, mask_1, .., mask_{K-1}]layoutvLLM runs at inference — and asserts the per-depth logits match the
cache_hiddenK-step training to ~1e-7. So the K sequential training stepsare numerically equivalent to vLLM's single parallel forward, not an
approximation. (This is what surfaced the RoPE bug fixed in #2377.)
Validation (8 GPUs, Llama-3.1-8B target, PerfectBlend, K=4)
End-to-end training. Loss decreases and accuracy rises — starting loss
≈
ln(draft_vocab = 8192) ≈ 9.0:The saved consolidated checkpoint carries the P-EAGLE additions in vLLM's
expected layout (
config.json:parallel_drafting=True,ptd_token_id;safetensors:
mask_hiddenof shape[1, 3*4096] = [1, 12288], plus thestandard EAGLE-3 draft keys). The acceptance-length / speedup numbers can only
be measured by serving the head under vLLM ≥ 0.16 on a GPU, which is out of
scope for this training-side PR.
Before your PR is "Ready for review"
Pre checks:
key/shape; masked-inputs-vs-no-recurrence;
num_draft_tokens=1== EAGLE-3step 0; loss-decreases; save round-trip config.json + safetensors;
out-of-draft-vocab masking; serve rejection)
Local:
ruffclean;pytest tests/unit_tests/speculative/test_peagle.py tests/unit_tests/speculative/test_eagle3.py tests/unit_tests/models/llama/test_rope_utils.py→ 54 passed, 2 skipped(GPU-only FlashAttention-2 paths).
Additional Information
support tracked in [Feature] Support P-EAGLE (Parallel EAGLE) speculative decoding sgl-project/sglang#23171.
verbatim with
Eagle3TrainerModule; factoring it into a common helper wouldtouch the well-tested TTT path and is left for a separate refactor.