Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training

AI summary

Overview

• Research area: Machine learning systems — specifically efficient reinforcement learning (RL) post-training for large language models, combining speculative decoding, distributed training parallelism, and online draft model adaptation.
• Technical level: Advanced. The paper assumes familiarity with speculative decoding, context parallelism (CP), pipeline parallelism (PP), ring attention, and GRPO-style RL post-training.
• One-sentence scope: The paper presents an end-to-end system that makes online co-training of target-feature-conditioned speculative draft models practical during large-scale, long-context RL post-training, by extending existing context- and pipeline-parallel layouts rather than changing them.

What This Paper Is About

The wall-clock time of RL post-training is dominated by rollout generation, and speculative decoding speeds that up by having a small draft model propose tokens that the large policy verifies in parallel. Because the policy keeps changing during RL, the draft must be co-trained online to keep acceptance length (the mean number of tokens accepted per verification) high.

Scaling that co-training to large models and long contexts breaks in two places: advanced drafts such as EAGLE-3, DFlash, and DSpark use branch-structured attention that standard causal context parallelism cannot express, and they consume intermediate hidden states from target layers that may sit on remote pipeline-parallel stages. The paper's goal is a system that supports both without disturbing the policy's existing parallel topology.

Key Contributions

1. Branch attention under context parallelism. The authors decompose draft branch attention into a causal main-sequence component, sharded across CP ranks, and a rank-local branch component kept on the rank owning the branch's anchor. The two attention results are merged with the standard online-softmax reduction (Eq. 2), within a packed, load-balanced zigzag-ring execution that supports EAGLE-3, DFlash, and DSpark.

2. Out-of-schedule target-feature transport under pipeline parallelism. They introduce TapChannel, a per-source mailbox on the draft stage with sequence stamps for producer–consumer handshake, delivering intermediate target features across PP stages on a side path that does not enter or alter the pipeline schedule. Cross-node sources use a dedicated NCCL communicator with GPUDirect RDMA; colocated sources use CUDA IPC.

3. End-to-end draft co-training at scale. Both mechanisms are integrated into the NeMo-RL framework, enabling online co-training of target-feature-conditioned drafts alongside an evolving policy under the joint objective L(θ, φ) = L_RL(θ) + λ·L_draft(φ, x, sg(H_θ(x))), with a stop-gradient on the target features.

4. Broad empirical validation. The system is evaluated across three draft families (EAGLE-3, DFlash, DSpark), target model sizes from 8B to 122B, single- and multi-turn RL workloads, plus operator-level CP benchmarks against USP and transport microbenchmarks.

Main Findings

• RL learning trajectory is preserved. On Qwen3-8B across EAGLE-3, DFlash, and DSpark, reward, AIME 2024 validation accuracy, and training–inference KL divergence closely track the baseline without speculative decoding. The paper states that a near-zero KL confirms the inference engine is correctly synchronized with the training policy.
• End-to-end results across scales (Table 1). Co-trained drafts reach 2.28–4.78 acceptance length, yielding 1.19–2.23× rollout speedup and 1.16–1.88× end-to-end training speedup. Per configuration: Qwen3-8B + EAGLE-3 = 2.28 / 1.63× / 1.50×; Qwen3-8B + DFlash = 3.45 / 2.23× / 1.88×; Qwen3-8B + DSpark = 3.63 / 2.18× / 1.83×; Qwen3.5-35B-A3B + DFlash = 4.58 / 1.50× / 1.46×; Nemotron-3.5-Lightning-30B-A3B + DSpark = 2.65 / 1.19× / 1.16×; Qwen3.5-122B-A10B + DFlash = 4.78 / 1.72× / 1.35×; GPT-OSS-120B + DFlash = 3.80 / 1.48× / 1.19×.
• Draft family matters. DFlash and DSpark consistently outperform EAGLE-3 in acceptance length. The larger MoE targets achieve high acceptance lengths but lower end-to-end speedups because each verification forward invokes sparse routing over more expert compute; the authors also note that for linear-attention models the verification step is comparatively cheaper, so relative speculation speedup is inherently smaller.
• Multi-turn workloads cost speedup. On the NeMo Gym Workplace Assistant with Qwen3-8B, all configurations achieve consistent rewards and acceptance length improves monotonically, but end-to-end speedup is 1.25–1.43× versus rollout-phase speedup of 1.75–2.23×, because rollout accounts for only 55.8% of step time and tool execution and environment latency are unreachable by faster decoding.
• CP beats USP. Against the USP implementation from SpecForge under matched GPU counts, on EAGLE-3 TTT (3 passes) with a longest sequence of 20,480 tokens, packed zigzag attention beats the best USP variant by 2.9×, 2.3×, and 1.5× in latency at CP=2, 4, and 8, with 2.7× lower per-GPU peak memory. USP pads the batch to 2.25× the real token count; the packed implementation avoids that.
• Long-context scaling. For TTT attention at fixed global lengths up to 256K tokens, latency drops from 17.7 s at CP=1 to 2.35 s at CP=8 (7.5×, 94% parallel efficiency), and per-GPU memory falls nearly linearly from 53.2 GB to 7.5 GB. Block-draft attention is launch-bound rather than FLOP-bound: it stays cheap even at CP=8 (≤ 55 ms from 16K to 256K) yet scales 6.9× at 256K when the workload is substantial.
• TapChannel transport. On a PP=4 fan-in with real tap layouts, one-sided writes achieve 27–39 GB/s and complete the fan-in 4.5–8.5× faster than staging through pinned host memory. Transfers introduce negligible interference on source stages and only 1.6% HBM contention on the receiving draft stage, whereas host staging slows every rank by 83–88% (also stated as >80%).
• PP overhead is modest (Table 2, Qwen3-8B, averaged over first 10 policy update steps). Draft co-training adds under 15% update-time overhead for block drafts and 34.3% for EAGLE-3 (due to its TTT passes), while speculative rollout cuts generation time by 28–52%, giving net speedups of 1.31–1.85×. The baseline row reports rollout time 329.2 and update time 31.5 (1.00×). EAGLE-3: accepted length 1.89, rollout time 238.7, 1.38× rollout speedup, update time 42.4, 34.3% overhead, tap wait 0.37, 1.31× E2E. DFlash: 2.84, 229.2, 1.44×, 35.8, 13.6%, 0.57, 1.37×. DSpark: 3.37, 156.8, 2.10×, 36.1, 14.6%, 0.54, 1.85×.
• Rendezvous cost is largely hidden. The draft stage waits only 0.4–0.6 s per policy update for taps, or 1.5–2.2% of optimization time, confirming that the cost is mostly overlapped by normal pipeline scheduling.

Methodology in Plain English

The authors start from an existing RL training system (NeMo-RL) that already uses CP and PP to spread a large policy across GPUs, and they refuse to change that topology. Instead they adapt the system around it.

For the attention problem, they observe that each draft branch query needs to look at two things: the long causal prefix of the main sequence, and a small set of keys belonging to its own branch. The first part behaves exactly like ordinary causal context parallelism, so it is computed with packed zigzag-ring attention where K/V circulate between ranks and queries stay put. The second part is small and lives entirely on the rank that owns the branch anchor, so it is computed locally with no extra communication. The two attention outputs are then combined using the same log-sum-exp merging trick that ring attention already uses between steps, so the result is mathematically equivalent to computing the full attention at once. This costs nothing extra in communication regardless of how many branches exist, and the backward pass reuses the same ring with no additional branch-dependent traffic. Communication is overlapped with computation at ring-step granularity.

For the feature-transport problem, target hidden states (taps) are produced on several PP stages while the draft lives only on the last stage, and normal pipeline communication only connects neighbors. Because taps need no return path, the authors give each source stage a pre-allocated mailbox slot in the draft stage's memory. After a source finishes its microbatch forward, it writes its taps; the draft stage reads the slot just before its own forward for that same microbatch. Sequence stamps coordinate the producer–consumer handshake, and separate mechanisms manage buffer reuse for local versus cross-node sources. Features are transferred directly once, no matter how many PP hops separate source and draft.

The evaluation trains with GRPO on DAPOMath-17K and validates on AIME 2024 (4,096 input / 16,384 response tokens on H100 GPUs), plus a multi-turn agentic workload on the NeMo Gym Workplace Assistant (based on WorkBench) with 32,768-token sequences. Larger models are run on GB200 GPUs.

Why This Matters

• Research impact: The paper reframes draft co-training as a systems problem rather than an objective-design problem. By keeping the target policy's parallel topology unchanged and adapting attention and feature routing around it, it removes a practical barrier that previously forced draft training onto separate or modified parallel layouts.
• Real-world applications:
  • Large-scale RL post-training of reasoning models, where rollout generation dominates wall-clock cost and even a 1.2–1.9× end-to-end speedup translates into substantial compute and time savings.
  • Agentic and tool-use training pipelines, where the paper's multi-turn results quantify how much of the achievable speedup is lost to tool execution and environment latency.
  • Long-context training and serving up to 256K tokens, where the CP decomposition delivers large memory and latency reductions relative to padded USP attention.
  • Deployment of target-feature-conditioned draft architectures (EAGLE-3, DFlash, DSpark) inside existing distributed training stacks without per-architecture layout rework.
• Industry relevance: The work is authored at NVIDIA and integrated into NeMo-RL, and its PP transport path explicitly targets GPUDirect RDMA and CUDA IPC. That places it directly in the tooling path for organizations training large MoE and dense models on multi-node H100 and GB200 clusters.

Future Directions

• Tailoring speculation for sparse MoE and linear-attention targets. The paper states this as explicit future work, motivated by its own finding that MoE targets and linear-attention models see limited relative benefit.
• Closing the multi-turn gap. Rollout was only 55.8% of step time in the Workplace Assistant workload; how to reduce or overlap tool-execution and environment latency is left open.
• Reducing EAGLE-3's co-training overhead. EAGLE-3's TTT passes drove a 34.3% update-time overhead versus under 15% for block drafts; whether that cost can be lowered is not addressed.
• Extending beyond the evaluated configurations. The paper reports no experiments beyond CP=8 or beyond the tested TP/PP/CP/EP combinations, so the behavior of the branch-attention decomposition at more extreme parallel degrees remains untested.

Target Audience

This paper benefits most readers working on distributed training systems and inference infrastructure for large language models — particularly engineers and researchers implementing speculative decoding inside RL post-training pipelines, and those responsible for context-parallel and pipeline-parallel training stacks at scale. It is also relevant to researchers studying draft model architectures (EAGLE-3, DFlash, DSpark) who need to understand the systems constraints their designs impose. It is not an introductory text: it presupposes comfort with ring attention, online-softmax merging, GRPO, and MoE parallelism notation.

Authors’ abstract

Speculative decoding accelerates rollout generation, which dominates the cost of reinforcement learning (RL) post-training. Online co-training can further increase the draft's accuracy, yielding greater speedups. However, scaling this approach to co-training on large models with long contexts poses two obstacles: (1) branch attention is unsupported by standard causal context-parallel (CP) implementations, and (2) target features span across pipeline-parallel (PP) stages. We address both with an end-to-end system for large-scale online draft co-training. For CP, we extend packed, load-balanced zigzag ring attention by merging rank-local branch attention with causal main-sequence attention. For PP, TapChannel transports intermediate target features across stages via a separate path, leaving the pipeline schedule unaffected. Experiments demonstrate that co-trained drafts closely track the policy baseline while delivering substantial rollout and end-to-end speedups across model scales up to 122B. Our CP design achieves strong scaling at 256K tokens with significant memory savings over prior work, and our PP transport incurs modest overhead. Code can be found at https://github.com/NVIDIA-NeMo/RL/issues/3698.