PipelineRL: インフライト重み更新による LLM 強化学習の最適化

PipelineRL は、大規模 LLM トレーニングにおいて高い推論スループットとオンポリシーデータ収集という根本的なトレードオフを解決することを目的とした実験的な強化学習 (RL) 実装です。インフライト重み更新 を実装することで、PipelineRL は推論サーバーが推論プロセスを停止せずに更新されたモデル重みを受け取れるようにし、高い GPU 利用率を確保しながらトレーニングデータをオンポリシーに近い状態に保ちます。

推論スループットとオンポリシーのトレードオフの解決

従来の RL ワークフローでは、効率性とデータの鮮度の間に矛盾があります。高いスループットを実現するために、推論サーバーは通常大きなバッチサイズを使用し、これにより複数のポリシー最適化ステップ用のデータが生成されます。しかし、各最適化ステップのたびに、データ収集に使用された重みと現在のポリシー重みとの「遅れ」が増大し、データはますますオフポリシーになり、トレーニングに対する効果が低下します。

PipelineRL は、推論を停止せずに各オプティマイザステップの後で推論サーバーの重みを更新することでこの問題を解決します。システムは新しい重みを受け取るために必要な短時間だけ推論サーバーを一時停止します。このアプローチにより、推論サーバーは最適なバッチサイズを維持しつつ、データがオンポリシーまたはオンポリシーに近い状態を保ち、より安定で効果的な学習を実現します。

パフォーマンスと安定性の結果

Open-Reasoner-Zero データセットでトレーニングした 7B および 32B モデルを用いた実験において、PipelineRL は AIME 2024 と MATH 500 の推論ベンチマークで Open-Reasoner と同等、またはそれ以上の性能を示すことが確認されました。

簡略化された RL アルゴリズム

競争力のある性能にもかかわらず、PipelineRL は Open-Reasoner-Zero よりもはるかにシンプルな RL 実装を使用しています。主な簡略化点は以下の通りです：

  • Simplified GRPO: 値関数を持たない、Group Relative Policy Optimization (GRPO) の簡略版を使用します。
  • No Complex Filtering: 実装ではトラストリージョン重要度重みのクランプ、過長シーケンスのフィルタリング、報酬シェイピングを省略しています。
  • Basic Loss Normalization: 損失はバッチ内のシーケンス数を分母として正規化し、すべてのトークンに同等の重みを与えます。
  • No Penalties: システムは KL ペナルティやエントロピー ボーナスを使用しません（ただし、リファレンスモデルの KL はサポートされています）。

KV キャッシュの古いデータの影響

インフライト重み更新における主な懸念は、シーケンス生成がモデルの以前のバージョンで計算された古いキーとバリューを KV キャッシュに保持したまま続行されることです。しかし、実験結果はこれがトレーニングの安定性に悪影響を与えないことを示しています。

モジュラーアーキテクチャと技術的契約

PipelineRL は、SGLang、vLLM などの専門的な推論ソフトウェアや、DeepSpeed、FSDP、TorchTitan などのトレーニングソフトウェアとの統合を可能にするためにモジュラー設計となっています。これは、主に 2 つの契約を通じて実現されます：

推論契約

PipelineRL と統合するために、推論ソフトウェアは以下の 3 つの特定 API を提供する必要があります：

  1. Process group initialization: 重み更新用のプロセスグループを初期化するための HTTP POST /init_process_group リクエスト。
  2. Weight Update Trigger: 推論サーバーに一時停止を指示し、NCCL を介して重みのブロードキャストを受け取るための HTTP POST /request_weight_update リクエスト。
  3. Chat completion: アクターとのやり取りのための標準的な HTTP POST /v1/chat/completion リクエスト。

トレーナー契約

トレーニングソフトウェアは、以下の操作のための Python API を提供する必要があります：

  • Worker initialization: トレーニング重みとオプティマイザ状態のロードおよびシャーディング。
  • Forward pass: トークンの対数尤度を生成。
  • Backward step: RL 目的関数の勾配を計算・蓄積。
  • Optimizer Step: オプティマイザステップを実行。
  • Weight gathering and broadcasting: 更新された重みを層ごとに収集し、推論サーバーへブロードキャスト。

実験設定

PipelineRL は、7B と 32B のモデルで以下のハイパーパラメータでテストされました：

  • Batch size: 4096
  • Learning rate: 1e-6
  • Max generated tokens: 8192

トレーニングに必要な計算リソースは、7B モデルで 2 ノードで約 3.5 日、32B モデルで 4 ノードで約 6 日でした。

Sources