Search papers, labs, and topics across Lattice.
To alleviate the compute-heavy rollout bottleneck in RL post-training, this work scales online speculative draft co-training to massive models and ultra-long contexts by solving distributed branch attention and activation routing challenges. The authors integrate rank-local branch attention into packed, load-balanced zigzag ring attention for Context Parallelism (CP) and introduce TapChannel to route intermediate target activations across Pipeline Parallelism (PP) stages without pipeline stalls. Evaluated up to 122B parameters and 256K context windows, the framework achieves strong scaling and substantial end-to-end rollout speedups while maintaining high speculative acceptance rates against the target policy.
Co-training speculative draft models directly inside 122B, 256K-token distributed RL runs removes the massive rollout bottleneck without causing pipeline stalls or context-parallel memory blowups.
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.
