NVIDIA-NeMo / NVIDIA-NeMo/RL

DFlash/DSpark and EAGLE-3 TTT draft co-training for speculative decoding (Megatron backend)

Open
#3,698 0 comments 0 reactions 1 assignee Claimed by @terrykong View on GitHub
community-request Feature Speed waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.