Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training
NVIDIA researchers detail an end-to-end system for online draft co-training that speeds speculative decoding in large-scale long-context RL post-training.
The paper tackles scaling online draft co-training for speculative decoding in RL post-training, where rollout generation dominates cost. It extends packed, load-balanced zigzag ring attention to merge rank-local branch attention with causal main-sequence attention for context parallelism, and introduces TapChannel to transport target features across pipeline-parallel stages without changing the schedule. Experiments show co-trained drafts tracking the policy baseline with substantial rollout and end-to-end speedups up to 122B parameters and strong scaling at 256K tokens.
- Extends packed zigzag ring attention with rank-local branch attention for causal context parallelism.
- TapChannel transports target features across pipeline-parallel stages without altering the schedule.
- Co-trained drafts track the policy baseline with substantial rollout and end-to-end speedups.
- Validated up to 122B parameters and 256K context with memory savings.
Full article160 words · extracted from huggingface.co · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://huggingface.co/papers/2609.07108