Speculative decoding is a key lever for accelerating rollout generation, which dominates the cost of reinforcement learning (RL) post-training. According to a new research paper, online co-training can further improve the draft model's accuracy and yield greater speedups — but scaling this to large models with long contexts has been difficult.
The authors identify two main obstacles. First, branch attention used by advanced draft models is not supported by standard causal context-parallel (CP) implementations. Second, the target features needed by the draft may reside on non-adjacent pipeline-parallel (PP) stages.
To address the CP issue, the system extends packed, load-balanced zigzag ring attention by decomposing branch attention into causal-sequence and rank-local components, merging rank-local branch attention with causal main-sequence attention. For the PP issue, a mechanism called TapChannel transports intermediate target features across stages via a separate path, leaving the pipeline schedule unaffected.
Experiments show that co-trained drafts closely track the policy baseline while delivering substantial rollout and end-to-end speedups across model scales up to 122B. The CP design reportedly achieves strong scaling at 256K tokens with significant memory savings over prior work, while the PP transport incurs only modest overhead.
The paper positions these contributions as an end-to-end system for large-scale online draft co-training in long-context RL post-training. Code is linked via the NVIDIA-NeMo/RL repository.