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

An end-to-end system enabling online draft co-training across context-parallel and pipeline-parallel stages for models up to 122B parameters

Zili Wang; Zhaopeng Qiu; Yuekai Zhang; Shuang Yu; Junjie Lai· 2026· DOI 10.48550/arXiv.2609.07108

The core problem

Reinforcement learning (RL) post-training of large language models is dominated by the cost of rollout generation. Speculative decoding is a well-established technique to accelerate this generation by using a smaller draft model to propose tokens that a larger target model verifies in parallel. While offline draft training is common, online co-training—where the draft is continuously updated alongside the target policy—can further increase draft accuracy and yield greater speedups. However, scaling online co-training to large models with long contexts introduces two fundamental obstacles: (1) branch attention is unsupported by standard causal context-parallel (CP) implementations, and (2) target features span across pipeline-parallel (PP) stages. This work addresses both challenges with an end-to-end system designed for large-scale online draft co-training. The authors 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 parameters. Their CP design achieves strong scaling at 256K tokens with significant memory savings over prior work, and their PP transport incurs only modest

Innovation

The authors conduct experiments across model scales up to 122B parameters and context lengths up to 256K tokens. Key findings include: - **Rollout and end-to-end speedups:** Co-trained drafts deliver substantial speedups in both rollout generation and end-to-end RL post-training. The drafts closely track the policy baseline, ensuring that the acceleration does not come at the cost of generation quality. - **Context parallelism scaling:** The proposed CP design achieves strong scaling at 256K tokens. Compared to prior work, it yields significant memory savings, enabling longer context training without exceeding memory limits. - **Pipeline parallelism overhead:** The TapChannel transport incurs only modest overhead, making it practical for large-scale deployment. - **Model scale:** The system is validated on models up to 122B parameters, demonstrating its applicability to state-of-the-art large language models. Quantitative results are not fully detailed in the abstract, but the authors emphasize that the co-trained drafts closely track the policy baseline while delivering substantial speedups. The code is available at the provided GitHub repository (https://github.com/NVIDIA-NeMo/

Reinforcement learning (RL) post-training of large language models is dominated by the cost of rollout generation. Speculative decoding is a well-established technique to accelerate this generation by using a smaller draft model to propose tokens that a larger target model verifies in parallel. While offline draft training is common, online co-training—where the draft is continuously updated alongside the target policy—can further increase draft accuracy and yield greater speedups. However, scaling online co-training to large models with long contexts introduces two fundamental obstacles: (1) branch attention is unsupported by standard causal context-parallel (CP) implementations, and (2) target features span across pipeline-parallel (PP) stages. This work addresses both challenges with an end-to-end system designed for large-scale online draft co-training. The authors 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 parameters. Their CP design achieves strong scaling at 256K tokens with significant memory savings over prior work, and their PP transport incurs only modest overhead.

The proposed system tackles the two obstacles through specialized parallelization strategies.

Why it matters

The work addresses a critical bottleneck in RL post-training: the cost of rollout generation. By enabling online draft co-training at scale, the system allows for continuous improvement of the draft model, leading to higher acceptance rates and greater speedups. The two technical contributions—branch attention in CP and TapChannel in PP—are essential for scaling to large models and long contexts. The CP design extends packed, load-balanced zigzag ring attention to support branch attention, which is necessary for speculative decoding. This is a non-trivial extension because branch attention introduces additional complexity in load balancing and causality. The memory savings at 256K tokens are particularly important, as long-context training is increasingly common. The PP transport via TapChannel is a clever solution to the problem of feature distribution across pipeline stages. By using a separate communication path, the authors avoid disrupting the pipeline schedule, which would otherwise lead to inefficiencies. The modest overhead suggests that this approach is practical for real-world deployment. One limitation is that the abstract does not provide detailed quantitative comparisons with prior work, such as exact speedup numbers or memory savings percentages. Future work could explore the trade-offs between draft model size and acceptance rate, as well as the impact of co-training frequency. Additionally, the system's performance on even larger models (e.g., 500B+ parameters) and longer contexts (e.g., 1M tokens) remains to be tested. Overall, this paper presents a significant step towards efficient RL post-training for large language models, with open-source code to facilitate further research.

Who should read this

CS practitioners and researchers

Opening member content…