用於大規模長上下文強化學習後訓練中推測解碼的線上草稿協同訓練

摘要

推測解碼可加速 rollout 生成，而 rollout 生成主導了強化學習（RL）後訓練的成本。線上共同訓練可進一步提升草稿模型的準確率，帶來更大的加速。然而，將此方法擴展至具長上下文的大型模型上進行共同訓練，會面臨兩項障礙：（1）標準因果上下文平行（CP）實作不支援分支注意力；（2）目標特徵跨越管線平行（PP）階段。我們以一套用於大規模線上草稿共同訓練的端到端系統解決這兩項問題。針對 CP，我們擴展打包、負載平衡的 zigzag 環狀注意力，將各 rank 本地的分支注意力與因果主序列注意力合併。針對 PP，TapChannel 透過獨立路徑跨階段傳輸中間目標特徵，且不影響管線排程。實驗顯示，共同訓練的草稿模型能緊密貼近策略基線，同時在模型規模高達 122B 的範圍內帶來顯著的 rollout 與端到端加速。我們的 CP 設計在 256K tokens 下達到強擴展性，並相較先前工作顯著節省記憶體；我們的 PP 傳輸則僅帶來適度的額外開銷。程式碼可在 https://github.com/NVIDIA-NeMo/RL/issues/3698 找到。

English

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.