Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training
Online Draft Co-Training with Zigzag Ring Attention and TapChannel achieves up to 1.88× end-to-end RL speedup.
Key Findings
Methodology
The paper introduces Online Draft Co-Training: policy parameters θ optimize the RL objective while draft parameters ϕ learn from the same rollout tokens and intermediate policy features Hθ under stop-gradient. The joint objective is L(θ,ϕ)=LRL(θ)+λLdraft(ϕ;x,sg(Hθ(x))). Without changing the policy’s TP/PP/CP topology, Zigzag Ring Attention handles causal main-sequence attention, branch KV remains rank-local, and TapChannel transports target features across pipeline stages.
Key Results
- On DAPOMath-17K with AIME 2024 evaluation, Qwen3-8B co-training with EAGLE-3, DFlash, and DSpark closely follows the no-draft RL baseline in reward, validation accuracy, and training–inference KL. Acceptance lengths are 2.28, 3.45, and 3.63, while end-to-end speedups are 1.50×, 1.88×, and 1.83×.
- Across targets from 8B to 122B, Qwen3.5-35B-A3B+DFlash reaches 1.46× end-to-end speedup, Qwen3.5-122B-A10B+DFlash 1.35×, and GPT-OSS-120B+DFlash 1.19×. Rollout speedups range from 1.19× to 2.23×.
- At 256K tokens, EAGLE-3 TTT latency falls from 17.7 seconds at CP=1 to 2.35 seconds at CP=8, giving 7.5× scaling and 94% parallel efficiency; memory drops from 53.2GB to 7.5GB. Compared with USP, latency improves by up to 2.9× and memory by 2.7×.
Significance
This work addresses the coupled bottlenecks of RL post-training: expensive rollout generation, policy–draft mismatch, and long-context memory and communication costs. It reframes speculative decoding as an online training component rather than merely an inference-serving optimization. Scientifically, it unifies branch attention, target-feature distillation, and large-scale distributed training. Industrially, it supports models from 8B to 122B and both single-turn and tool-using multi-turn workloads, converting rollout improvements into 1.16–1.88× end-to-end training acceleration.
Technical Contribution
The first contribution is a branch-attention decomposition: main-sequence causal attention and rank-local branch attention are computed separately and merged with online softmax, ℓ=log(eℓm+eℓb) and O=eℓm−ℓOm+eℓb−ℓOb. This supports EAGLE-3, DFlash, and DSpark. The second is TapChannel, an out-of-schedule side path using CUDA IPC or GPUDirect RDMA, mailboxes, sequence stamps, and preallocated slots to fan in non-adjacent target features without modifying the pipeline schedule.
Novelty
The novelty is not the invention of speculative decoding or online draft training, but the system-level integration of multiple branch-structured drafts into an existing large-model CP/PP topology. Unlike SpecForge’s sequential ring and USP’s padding-oriented implementation, the method uses packed, load-balanced zigzag-ring attention. Unlike standard adjacent-stage pipeline communication, TapChannel provides one-way, non-adjacent, schedule-independent feature transport.
Limitations
- End-to-end gains depend on how much of a training step is attributable to rollout. In Workplace Assistant, tool execution and environment latency are not accelerated by decoding, limiting total speedup to roughly 1.25–1.43×.
- Strong scaling shortens each rank’s local context and eventually exposes communication as a bottleneck. For MoE targets, verification still invokes sparse expert computation, so high acceptance length does not guarantee proportional training acceleration.
- The evaluation is concentrated on NeMo-RL, selected models, and official draft checkpoints; broader reward functions, hardware systems, and draft objectives remain underexplored.
Future Work
Future work should improve fine-grained communication–attention overlap, dynamically choose draft lengths and architectures, and optimize directly for acceptance. Important directions include longer multi-turn trajectories, heterogeneous clusters, alternative MoE routing policies, larger targets, and treating TapChannel as a general cross-stage feature service. More systematic studies of λ, draft losses, and policy–draft co-adaptation are also needed.
AI Executive Summary
Rollout generation often dominates the wall-clock cost of reinforcement-learning post-training. Speculative decoding lets a small draft model propose several tokens before the target policy verifies them in parallel, but a fixed draft drifts as the policy evolves. Scaling online co-training is harder still: EAGLE-3, DFlash, and DSpark require branch-structured attention and intermediate target features, whereas standard causal context parallelism and pipeline parallelism support neither naturally.
Wang and colleagues propose Online Draft Co-Training, jointly updating the policy and draft while stopping gradients through target features. Zigzag Ring Attention splits each branch query into causal main-sequence attention and rank-local branch attention, then merges them with an online-softmax reduction. TapChannel sends non-adjacent target features through a separate path to the draft stage, leaving the original pipeline schedule intact. The design supports all three draft families.
On DAPOMath-17K with AIME 2024 evaluation, Qwen3-8B plus DFlash achieves acceptance length 3.45, 2.23× rollout speedup, and 1.88× end-to-end speedup. Across 8B–122B targets, end-to-end gains range from 1.16× to 1.88×. At 256K context, EAGLE-3 TTT reaches 94% parallel efficiency while memory falls from 53.2GB to 7.5GB. Tool latency and MoE verification remain limiting factors, but the work demonstrates that online drafting can become a practical component of large-scale RL infrastructure.
Deep Analysis
Background
Speculative decoding, introduced by Leviathan et al., uses draft proposals and parallel target verification while rejection sampling preserves the target distribution. Systems such as NeMo-RL integrate MTP and EAGLE-3; FastGRPO and ReSpec adapt drafts online. However, EAGLE-3 TTT, DFlash block diffusion, and DSpark’s Markov head introduce branch KV and target hidden-state interfaces. Conventional RingAttention handles causal main sequences but not these branch and cross-stage requirements.
Core Problem
The goal is to co-train drafts during large-model, long-context RL while preserving the policy’s existing TP/PP/CP layout. Under CP, a branch query must attend both to a sharded causal prefix and to local branch KV. Under PP, target features may be produced on non-adjacent stages while the draft resides on the last stage. Redesigning topology would increase memory, scheduling complexity, and engineering risk.
Innovation
First, packed Zigzag Ring Attention shards variable-length main sequences with load balance; main-sequence K/V circulates while branch K/V remains on its anchor rank. Second, online softmax merges the two attention components, making one operator compatible with EAGLE-3, DFlash, and DSpark. Third, TapChannel uses preallocated mailboxes, sequence stamps, CUDA IPC, and GPUDirect RDMA for schedule-independent feature fan-in. Unlike USP, it avoids 2.25× padding of the real token count.
Methodology
- �� Joint objective: train policy and draft on the same rollout tokens, feeding the draft stop-gradient target features so its auxiliary loss does not update policy representations.
- �� CP execution: branch queries remain resident on their owner rank; main-sequence K/V circulate for C ring steps. Compute (Om,ℓm) and (Ob,ℓb), then merge with ℓ=log(eℓm+eℓb) and O=eℓm−ℓOm+eℓb−ℓOb.
- �� Communication: prefetch the next K/V shard while computing the current attention step; exposed cost is max(tcomm,tattn). Forward outbound volume is 2(C−1)N/C·dkv·b.
- �� PP transport: source stages write taps into preallocated slots on the draft stage; sequence stamps coordinate producer and consumer. Cross-node transfers use NCCL plus GPUDirect RDMA, while colocated sources use CUDA IPC.
- �� Evaluation: joint GRPO training in NeMo-RL reports reward, AIME accuracy, training–inference KL, acceptance length, throughput, and total step time.
Experiments
Training uses DAPOMath-17K and AIME 2024 validation; multi-turn testing uses NeMo Gym Workplace Assistant at sequence length 32,768. Targets include Qwen3-8B, Qwen3.5-35B-A3B, Qwen3.5-122B-A10B, Nemotron-3.5-Lightning-30B-A3B, and GPT-OSS-120B. Drafts are EAGLE-3, DFlash, and DSpark. CP is compared with USP through 256K tokens, while PP overhead is compared with host-staged feature transfer.
Results
All three drafts preserve a learning trajectory close to baseline. On Qwen3-8B, DFlash reaches 1.88× end-to-end speedup and DSpark 1.83×; DFlash still reaches 1.35× on the 122B Qwen target. Against USP, packed zigzag attention is 2.9×, 2.3×, and 1.5× faster at CP=2, 4, and 8, respectively, with 2.7× lower per-GPU memory. TapChannel completes fan-in 4.5–8.5× faster than host staging and causes at most 1.6% HBM contention on the receiving stage.
Applications
The system fits mathematical reasoning, code generation, agentic tool use, and long-document RL post-training. Users need a distributed GPU cluster, NeMo-RL-compatible policy training, target intermediate-feature access, and a trainable draft module. For online services, continuously adapted drafts reduce policy drift; for training platforms, existing TP/PP/CP topology can remain unchanged.
Limitations & Outlook
Speedup is not determined by acceptance length alone: tool execution and environment waiting are unreachable by speculative decoding, while MoE verification can require substantial expert computation. Extreme CP scaling exposes communication bottlenecks, and TapChannel must manage cross-node buffers and in-flight transfers. The paper does not provide exhaustive ablations over draft losses, λ, rewards, hardware, or alternative co-training objectives.
Plain Language Accessible to non-experts
Imagine a large restaurant whose head chef represents the target model. Cooking every dish from scratch is slow, so a junior assistant prepares several likely dishes first. The chef checks them together; approved dishes reach customers quickly, while rejected ones are corrected according to the chef’s own recipe. Because the chef still makes the final decision, the restaurant’s quality standard does not change.
Now imagine that the restaurant has many floors and a very long order split across workstations. Ordinary service elevators connect only neighboring floors and do not understand temporary side-orders. The paper creates two special systems. Zigzag Ring lets the main order information circulate among workstations while keeping each side-order locally. TapChannel acts like a private express elevator, carrying important ingredients from distant kitchens directly to the assistant’s workstation without disturbing the normal production schedule.
With Qwen3-8B and DFlash, rollout becomes 2.23× faster and total training 1.88× faster. Even a 256K-token order can be handled efficiently. But waiting for customers, tools, or complicated expert kitchens still takes time, so the total restaurant improvement is smaller than the cooking improvement.
ELI14 Explained like you're 14
Imagine a huge video game where the main player is brilliant but takes a long time to decide every move. A smaller sidekick guesses the next few moves first. The main player checks the guesses all at once, so correct guesses save loads of time. If a guess is wrong, the main player fixes it, meaning the final game behavior stays trustworthy!
The tricky part is that the main player keeps learning. An old sidekick may stop guessing correctly, so the paper trains the sidekick at the same time as the main player. EAGLE-3 predicts step by step, DFlash guesses a whole block, and DSpark gives that block a lightweight left-to-right correction.
The computer may spread the giant game world across many GPUs. Zigzag Ring is like passing the main map around a team while letting each player keep their local mini-map. TapChannel is a private delivery lane that sends useful information from a faraway GPU to the sidekick without messing up the main workflow.
The results are exciting: Qwen3-8B with DFlash gets 2.23× faster rollouts and 1.88× faster total training. It even scales to 256K tokens! But if the game must wait for tool responses, or every check needs lots of expert work, the whole game cannot speed up by the same amount.
Glossary
Speculative Decoding
A small model proposes several tokens and a larger target model verifies them in parallel. Rejection sampling preserves the target model’s output distribution.
It accelerates RL rollout generation in this paper.
Online Draft Co-Training
The draft is continuously trained alongside an evolving RL policy. It consumes policy features through stop-gradient so draft learning does not backpropagate into the policy.
This is the paper’s central training framework.
Context Parallelism, CP
A long sequence is distributed across GPUs, with K/V communication enabling global attention. This reduces per-GPU memory for long contexts.
The paper extends CP to branch attention.
Pipeline Parallelism, PP
Model layers are partitioned into stages and microbatches flow through them. Standard communication usually connects only adjacent stages.
TapChannel overcomes this restriction.
TapChannel
An out-of-schedule side path for transporting target features. It uses mailboxes, sequence stamps, CUDA IPC, or GPUDirect RDMA.
It delivers remote features to the final-stage draft module.
Acceptance Length
The average number of draft tokens accepted during one target verification. Higher values generally imply more effective speculative decoding.
Reported values range from 2.28 to 4.78.
Open Questions Unanswered questions from this research
- 1 The effects of draft loss design, λ, and reward structure on acceptance and policy stability are not systematically disentangled.
- 2 It remains unclear whether communication, buffering, and scheduling dominate beyond 256K tokens or on more heterogeneous clusters.
- 3 Joint optimization of tool latency, MoE routing, and draft acceptance remains unresolved.
Applications
Immediate Applications
Long-context mathematical RL
Research teams can integrate DFlash or DSpark into NeMo-RL and GRPO for DAPOMath-like training. A CP/PP-capable GPU cluster and target-feature interface are required; the paper suggests roughly 1.16–1.88× end-to-end acceleration.
Multi-turn agents and tool use
Agent systems can use EAGLE-3, DFlash, or DSpark to accelerate generation while preserving target-policy behavior. Tool and environment waiting should be measured separately because they cap total end-to-end gains.
Long-term Vision
Unified distributed speculative-training infrastructure
TapChannel and branch attention could become reusable infrastructure that automatically adapts diverse drafts to TP, PP, CP, and MoE systems, enabling longer trajectories and larger policies without redesigning parallel topology.
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.