大規模長上下文 RL 後訓練推測性解碼的線上草稿共同訓練
為何重要
此項研究針對大型模型 RL 後訓練（RLHF）中最耗時的生成環節提出系統級最佳化方案，透過解決 CP 與 PP 並行架構的限制，為開發者提供在不增加硬體負擔下提升訓練速度與質量的工具。此技術若能成熟整合，將增強 NVIDIA 等生態系統在效率上的競爭壁壘，並對算力供應鏈的商業化應用有直接意義。
推測性解碼大幅加速強化學習後訓練，但其在大規模長上下文模型中的線上共同訓練面臨上下文不支援分支注意力、以及目標特徵跨越管線並行階段的技術瓶頸。
- 為解決上下文並行問題，系統將編組負載平衡的 zigzag ring attention 與 rank-local branch attention 結合。
- 為解決管線並行問題，透過 TapChannel 在獨立路徑傳輸中間目標特徵，不影響原有管線排程。
- 實驗顯示在 122B 模型上，共同訓練草稿能緊密追蹤策略基準，並帶來顯著的 rollout 與整體啟動加速。
- 256K tokens 下具備強擴充套件性，相較先前工作節省大量記憶體，且管線開銷極小。