ServiceNow/PipelineRL

A scalable asynchronous reinforcement learning implementation with in-flight weight updates.

何を解決するか

PipelineRLは、大規模言語モデル(LLM)における強化学習(RL)の効率性のトレードオフを解決します。具体的には、高い推論スループット(多くのGPUにわたる大規模バッチが必要)を達成することと、"オンポリシー"データの新鮮さ(モデルが最新バージョンで生成したデータに基づいて学習されること)を維持することの間の矛盾を解消します。

動作方法

PipelineRLは、"飛行中の重み更新"を用いたスケーラブルな非同期アーキテクチャを使用します。モデルの更新のためにサンプリングプロセス全体を停止するのではなく、最適化ステップの直後にNCCLを介して更新された重みが推論サーバーに即座にブロードキャストされます。これにより、ポリシーの更新中にロールアウトの生成を継続でき、GPUの利用効率を損なうことなくデータをオンポリシーに近づけられます。

システムは6つのモジュール構成で構成されています:

  1. オーケストレーター:GPUの配置を管理し、サブプロセスを起動します。
  2. 推論サーバー:vLLMベースのサーバーで、サンプリングを処理し、重みの更新を受け取ります。
  3. アクタープロセス:LLMからサンプリングしてロールアウトを生成し、報酬を収集します。
  4. プリプロセッサー:シーケンスをトークン化し、アドバンテージを計算します。
  5. トレーナー:RLステップ(簡略化されたGRPOアルゴリズムを使用)を実行し、重みの更新をトリガーします。
  6. 検証者:オプションのサーバーで、モデル出力の正しさを検証します(例:数学タスク)。

対象ユーザー

大規模言語モデルエージェントの訓練を行う研究者や開発者向けです。特に、報酬が検証可能な推論タスク(数学やコーディングなど)に注力し、複数のGPUにわたって訓練をスケーリングする必要があるユーザーに適しています。

特徴

  • 飛行中の重み更新:サンプリングパイプラインを停止せずにモデルパラメータを更新します。
  • エージェント非依存load_problems および generate_rollout 関数を実装することで、任意のエージェントタスクに適応可能です。
  • 高いパフォーマンス:AIME-2024およびMATH-500ベンチマークでOpen-Reasoner-Zeroと同等またはそれを上回ります。
  • 柔軟なバックエンド:プロセス間通信にファイルシステムまたはRedisベースのストリーミングの両方をサポートします。
  • サンドボックス統合:SandboxFusionをサポートし、リモートサンドボックスでコードを実行・検証できます。

関連

  • Dispatch
  • プロジェクト
  • プロジェクト
  • Dispatch
  • プロジェクト