DFlash/DSpark and EAGLE-3 TTT draft co-training for speculative decoding (Megatron backend)
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Motivation
NeMo-RL already ships a single-pass EAGLE-3 draft-training path (`policy.draft`). When serving RL rollouts with speculative decoding, a *static* drafter's acceptance rate degrades as the policy trains away from it — co-training the drafter alongside the policy (and refitting it into vLLM together with the policy weights) keeps acceptance high and rollout generation fast.
## What I'm working on
Extending the draft co-training path in five directions:
1. **Multi-pass TTT training for EAGLE-3** (`draft.ttt_steps > 1`): sequential draft passes with a two-part trunk+branch flash attention (exact joint-LSE merge, inference-aligned per-pass RoPE offsets), instead of the current single pass.
2. **DFlash/DSpark block drafters** (`draft.speculator_type: dflash|dspark`): anchor + mask-token block drafting (bidirectional in-block attention, trunk truncated at the anchor), following the official DFlash contract — the draft owns no lm_head / mask embedding and rides the target's live head, matching vLLM ≥ 0.26 serving.
3. **Training-loop integration**: chunked vocab-parallel soft-CE distillation losses with per-pass/per-slot global normalization, a separate draft grad-norm clipping group, an optional draft-only optimizer param group (FastGRPO-style higher LR), and vLLM refit fixes (drafter lm_head un-sharing, spec-decode max_tokens clamping).
4. **Context parallelism + sequence packing for the draft** (`sequence_packing.enabled` with `megatron_cfg.context_parallel_size > 1`): the TTT and block-draft attentions run on the THD varlen layout with the pass-1 trunk K/V on a zigzag ring (exact joint-LSE merge), the draft losses gather in rank-local packed coordinates, sliding-window block layers all-gather their trunk K/V, and one `DraftCrossEntropyLossFn` covers single/multi-pass, packed/unpacked and CP.
5. **Pipeline parallelism for the draft** (`megatron_cfg.pipeline_model_parallel_size > 1`): the draft stays whole on the last stage; a one-sided tap channel (CUDA IPC intra-node, an NCCL side channel across nodes) delivers the aux hidden states and pass-1 embeddings per microbatch without touching the 1F1B schedule, plus every-rank grad-norm-group registration and a PP-broadcast refit of the draft weights.
Items 4–5 are the training-system side of our paper, https://arxiv.org/abs/2609.07108 (co-training drafters with the policy under CP/PP); this issue is the roadmap for upstreaming it.
## Scope
This covers the **Megatron** policy backend only — the automodel draft-training path is being worked on separately by others.
## Status
Split into a stacked series of five draft PRs (each PR's diff includes the ones below it until they merge):
1. #3699 — multi-pass TTT training for EAGLE-3 drafts (+ shared infra: chunked distillation CE, draft-only optimizer group, vLLM refit safety). Runs on the currently pinned vLLM.
2. #3712 — DFlash block-draft co-training (vLLM ≥ 0.25 serving).
3. #3713 — DSpark block-draft co-training with the official CE+TV+confidence loss (vLLM ≥ 0.26 serving, `drafter` → `speculator` rename tolerated).
4. #4059 — context parallelism + sequence packing for draft co-training (incl. sliding-window block layers under CP).
5. #4060 — pipeline parallelism for draft co-training via a one-sided tap channel.
Unit tests included throughout. Items 1–3 were smoke-tested on a Qwen3-8B GRPO math recipe with eagle3 and dflash drafters; items 4–5 are the code path the paper's experiments ran on — a wide range of configurations across drafter types, policy scales (up to 122B) and context lengths (up to 256K), not reproduced here; see the report for details.
Related finding, tracked separately: #3716 (split-API train step misses the TP grad all-reduce for replicated params).
Opening this issue for visibility — feedback on the direction is welcome, and happy to align with any related internal work.
Contributor guide
Assessment
This issue has not been assessed yet.