Source-linked AI summary
Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training
Zili Wang, Zhaopeng Qiu, Yuekai Zhang, Shuang Yu, Junjie Lai
TL;DR
Large-scale online draft co-training for speculative decoding must support branch attention under context parallelism and target features spanning pipeline stages. The paper introduces packed zigzag-ring attention and TapChannel to address these obstacles, and reports stable learning with substantial speedups across long-context workloads and model scales.
Problem
Online draft co-training can improve speculative-decoding speedups, but standard causal CP does not support branch attention and PP can place required target features on remote stages.
Method
The system merges causal and rank-local branch attention through packed zigzag-ring execution and transports intermediate target features across PP stages with TapChannel outside the pipeline schedule.
Results
Co-trained drafts closely track the policy baseline while delivering 1.50–1.88× end-to-end speedups across target scales and RL workloads.
Takeaways & Limitations
The system enables online co-training of EAGLE-3, DFlash, and DSpark for long-context RL post-training while preserving the target’s learning trajectory and achieving scalable CP performance.
Takeaways & Limitations
Aggressive strong scaling can eventually expose communication as the bottleneck as local context shrinks.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper develops an end-to-end system for online speculative-draft co-training in large-model, long-context RL by adapting context and pipeline parallelism. Experiments show stable learning and substantial speedups across model scales and workloads.
- Motivation: Rollout generation often dominates RL post-training cost, motivating speculative decoding and online draft co-training as acceleration strategies.Speculative decoding verifies multiple draft tokens in parallel, while online co-training adapts the draft as the RL policy evolves.
- Challenges: Large-scale co-training must handle branch attention unsupported by standard causal CP and target features distributed across remote PP stages.These constraints prevent the draft from simply inheriting the policy’s existing parallelism configuration.
- Results: 1.50–1.88× end-to-end speedups accompany co-trained drafts that closely track the baseline in reward and accuracy across scales and RL workloads.The evaluation covers three draft families, target models from 8B to 122B, and single- and multi-turn workloads.
- Method: Packed zigzag-ring attention merges causal main-sequence attention with rank-local branch attention to support EAGLE-3, DFlash, and DSpark.The mechanism decomposes branch attention into causal and local components within a packed, load-balanced execution.
- Method: TapChannel transports intermediate target features across pipeline stages on a side path without entering or altering the pipeline schedule.This extends the existing PP configuration while leaving the pipeline schedule unchanged.
2 RELATED WORK
Prior work accelerates speculative decoding and adapts drafts, but integrating target-feature-conditioned drafts into existing distributed systems remains difficult. This paper instead keeps the target topology unchanged while adapting the surrounding system.
- Speculative decoding: Speculative-decoding research has primarily optimized draft sourcing, serving configuration, and rollout-side verification.The paper positions continuous draft training inside the distributed policy learner as complementary to these efforts.
- Draft adaptation: Online draft adaptation addresses policy evolution by recovering acceptance length and speedup during RL training.Related approaches include online draft learning, dynamic speculative-parameter selection, policy distillation, and parameter-sharing MTP modules.
- Draft architectures: Target-feature-conditioned architectures trade proposal quality against sequential cost through mechanisms such as fused target features, recursive conditioning, and parallel candidate generation.The cited systems include EAGLE-3, FastMTP, and DFlash.
- Distributed integration: Existing distributed-training methods support causal attention but do not directly handle draft-specific branch attention or cross-stage feature routing.The paper distinguishes its setting from work that changes the parallel layout to accommodate draft training.
- Positioning: The paper keeps the target’s topology unchanged and adapts the surrounding system for draft training.This contrasts with previous work that modifies the parallel layout.
3 METHOD
The system enables online draft co-training for large-model, long-context RL by extending context and pipeline parallelism to handle branch attention and cross-stage target features. It combines branch-local and causal attention under CP with TapChannel feature transport under PP, while preserving pipeline scheduling and overlapping communication where possible.
- 3 METHOD: Online co-training supports large-model, long-context RL through coordinated extensions to context and pipeline parallelism.The system integrates mechanisms for branch attention and target-feature transport into the existing distributed training configuration.
- 3.2 Branch Attention under Context Parallelism: Branch attention is decomposed into causal main-sequence and rank-local branch components, then merged within packed zigzag-ring attention.Main-sequence keys and values circulate across CP ranks while branch-local computation remains on the owning rank.
- 3.2 Branch Attention under Context Parallelism: The CP design supports EAGLE-3, DFlash, and DSpark despite their different branch structures.The method accommodates EAGLE-3’s multi-step branches and block-oriented draft architectures within the same attention mechanism.
- 3.3 TapChannel under Pipeline Parallelism: TapChannel transports intermediate target features from non-adjacent PP stages to the draft stage through an out-of-schedule side path.Per-source mailbox slots and sequence stamps coordinate feature writes and reads without interfering with the pipeline schedule.
- 3.3 TapChannel under Pipeline Parallelism: Tap transfers add no latency when they complete within pipeline slack; only transfer time beyond that slack creates visible rendezvous delay.The design transfers each feature directly once, while the pipeline schedule provides the opportunity to hide communication.
- 3.4 Evaluation: Experiments evaluate reward, accuracy, KL divergence, accepted length, throughput, and end-to-end step time across online co-training runs.The evaluation includes Qwen3-8B experiments with EAGLE-3, DFlash, and DSpark on DAPOMath-17K.
4 EXPERIMENTS
Experiments evaluate learning fidelity, end-to-end acceleration, context-parallel scaling, and pipeline-parallel overhead across draft families, model scales, and workloads. The system preserves the RL learning trajectory while improving rollout and training efficiency, including at 256K-token contexts.
- Experimental Setup: Experiments cover three draft families, targets from 8B to 122B, single- and multi-turn workloads, and learning, speed, and overhead metrics.The evaluation uses EAGLE-3, DFlash, and DSpark with Qwen3-8B and larger targets, including Qwen3.5-122B-A10B and GPT-OSS-120B.
- Policy Learning: Co-trained runs closely track the baseline reward, validation, and training-inference KL learning trajectory across all three draft models.The comparison includes a baseline without online co-training or speculative decoding and co-trained EAGLE-3, DFlash, and DSpark runs.
- End-to-End Performance: Co-trained drafts reach 2.28–4.78 acceptance length, producing 1.19–2.23× rollout and 1.16–1.88× end-to-end training speedups across targets and draft families.DFlash and DSpark consistently outperform EAGLE-3 in acceptance length; larger MoE targets have lower end-to-end speedups because verification invokes more sparse expert computation.
- End-to-End Performance: Multi-turn Workplace Assistant runs achieve 1.25–1.43× end-to-end speedup versus 1.75–2.23× rollout speedup because rollout represents only 55.8% of step time.Tool execution and environment latency are inside the rollout phase but are not accelerated by faster decoding.
- Context-Parallel Performance: Packed zigzag attention beats the best USP variant by 2.9×, 2.3×, and 1.5× latency at CP=2, 4, and 8, with 2.7× lower per-GPU peak memory.The comparison uses the same variable-length workload, where USP pads to 2.25× the real token count and packed attention avoids that overhead.
- Context-Parallel Performance: At 256K tokens, TTT attention latency falls from 17.7 s at CP=1 to 2.35 s at CP=8, while per-GPU memory decreases from 53.2 GB to 7.5 GB.This corresponds to 7.5× scaling and 94% parallel efficiency; block-draft attention remains at or below 55 ms from 16K to 256K.
- Pipeline-Parallel Performance: TapChannel adds modest pipeline overhead while speculative rollout reduces generation time by 28–52%, yielding net speedups of 1.31–1.85×.Draft co-training overhead is under 15% for block drafts and 34% for EAGLE-3; tap waits consume 1.5–2.2% of optimization time.
5 CONCLUSION
The paper presents a system that enables online co-training of target-feature-conditioned speculative drafts under context and pipeline parallelism for long-context RL post-training. Experiments show preserved RL learning trajectories, consistent speedups across target scales, scalable long-sequence context parallelism with lower memory use, and modest pipeline-parallel overhead.
- The system supports online co-training of EAGLE-3, DFlash, and DSpark under context and pipeline parallelism for long-context RL post-training.
- All three drafts preserve the RL learning trajectory while delivering consistent end-to-end speedups across target scales.
- Packed zigzag-ring attention scales to long sequences with lower memory usage.
- TapChannel provides cross-stage feature transport with modest pipeline-parallel overhead.