Summary
Generating actions in trained AI models to improve them can take a lot of time and computer power. The authors explore a method called speculative decoding, which helps make this process faster by making good guesses ahead of time. They improve this with a technique called online co-training, which makes these guesses more accurate, especially for very large models with long sequences of information. To handle the technical challenges in big models, they designed a system that efficiently manages attention across different parts of the model and communicates needed data during training. Their experiments show that this approach keeps the model's performance close to the usual method but completes tasks faster and uses less memory.
Speculative decodingReinforcement learningOnline co-trainingCausal context-parallel attentionPipeline parallelismLarge-scale modelsRollout generationBranch attentionMemory efficiencyModel scaling
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.