Online co-training speeds up large models in reinforcement learning
Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training
Machine LearningDistributed, Parallel, and Cluster Computing
Summary
Generating the steps for training AI with reinforcement learning can be very slow, especially with large models and long sequences. The authors developed a new system that improves a technique called speculative decoding by training two drafts together online. Their system solves technical challenges so it works efficiently with very big models and long contexts, making the training faster without losing accuracy. This helps large-scale AI training run more quickly and use less memory.
What this means in practice
- •For machine learning engineers: Speed up reinforcement learning model training on very large models with long input contexts using improved speculative decoding.
- •For cloud infrastructure teams: Reduce memory usage and computation time for large-scale language model inference pipelines with long sequences and pipeline parallelism.
Authors
Zili Wang, Zhaopeng Qiu, Yuekai Zhang, Shuang Yu, Junjie Lai
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.