大規模・長文脈RLポストトレーニングにおける投機的デコーディングのためのオンライン・ドラフト共訓練

要旨

投機的デコーディングは、強化学習（RL）ポストトレーニングのコストの大部分を占めるロールアウト生成を高速化する。オンライン共訓練はドラフトの精度をさらに高め、より大きな高速化をもたらす。しかし、このアプローチを長いコンテキストを持つ大規模モデルでの共訓練にスケールさせるには、2つの障害がある。(1) 分岐アテンションは標準的な因果的コンテキスト並列（CP）実装ではサポートされていない。(2) ターゲット特徴量がパイプライン並列（PP）ステージ間にまたがる。我々は、大規模オンライン・ドラフト共訓練のためのエンドツーエンドシステムで両方に対処する。CPについては、パック済みで負荷分散されたジグザグ・リングアテンションを拡張し、ランクローカルな分岐アテンションと因果的主系列アテンションを統合する。PPについては、TapChannelが中間ターゲット特徴量を別経路でステージ間転送し、パイプラインスケジュールに影響を与えない。実験により、共訓練されたドラフトはポリシーベースラインに密接に追従しつつ、最大122Bまでのモデル規模で大幅なロールアウトおよびエンドツーエンドの高速化を実現することが示された。我々のCP設計は256Kトークンでストロングスケーリングを達成し、先行研究と比べて大幅なメモリ節約を実現する。また、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.