vLLM V1 移行：強化学習におけるバックエンドの正確性の確保

ServiceNow AI は、目的側の修正よりもバックエンドの正確性を優先することで、RL 推論エンジンを vLLM V0 から V1 へ正常に移行しました。チームは、処理されたロールアウト logprob、V1 固有のランタイムデフォルト、インフライトの重み更新パス、そして lm_head 投影の精度という 4 つの重要な不一致を解決することで、vLLM 0.8.5 (V0) と vLLM 0.18.1 (V1) のパリティを達成しました。

移行の目的

PipelineRL のようなシステムでロールアウト生成の推論エンジンとして vLLM を使用する場合、エンジンはトークンをサンプリングし、トレーナーがポリシー比率、KL ダイバージェンス、クリップレート、エントロピーなどの重要な RL 指標を計算するために使用する logprob を返します。これらの logprob の計算方法に差異があると、トレーニングと推論の不一致が生じ、トレーニングのダイナミクスを根本的に変えてしまいます。

vLLM V1 へ移行するために、ServiceNow AI は狭い目的を設定しました：V1 がトレーナーが期待する形式でロールアウト logprob を返すことを検証し、V0 参照に対してワークロードを再実行し、バックエンドのパリティが回復した後にのみ目的レベルの変更を評価することです。V1 での初期試行では、報酬、エントロピー、クリップレートが V0 参照と比較して大きくずれていることが示されました。

バックエンドの失敗モード

チームは、トレーニングと推論の不一致の潜在的原因を 3 つの層に分類しました：

  1. セマンティック不一致：バックエンドがトレーナーの期待する意味と異なる logprob を返す。
  2. 推論パスの不一致：キャッシュ、スケジューリング、リクエスト処理のランタイムデフォルトの違いにより、プロンプトが異なる実行パスをたどる。
  3. 目的の不一致：RL の目的が、残存する古さやバックエンドの不一致を補正する必要がある。

最初の 2 つをバックエンドの挙動問題として扱い、まずそれらを排除することで、チームはバックエンドエラーを隠すために目的側の修正を早期に適用するという誤りを回避しました。

vLLM V1 バックエンド修正

Logprob のセマンティクス

デフォルトでは、vLLM V1 はロジットの後処理（温度スケーリングや top‑k／top‑p フィルタリングなど）を行う前の生モデル出力から logprob を返します。PipelineRL はサンプラーが使用する処理済み分布からの logprob を必要とします。logprobs-mode=processed_logprobs を設定することで、ロールアウト logprob の平均オフセットが除去され、平均ポリシー比率が 1.0 に近づきました。

ランタイムデフォルト

パリティを確保するために、チームは V0 参照パスと異なる V1 固有のデフォルトを明示的に無効化しました：

  • Prefix Caching：enable-prefix-caching: false に設定し、重み更新前に計算された状態の再利用を防止し、キャッシュ寿命に V1 のみの差異が生じるのを防ぎました。
  • Async Scheduling：async-scheduling: false に設定し、別の V1 のみの自由度を排除しました。

インフライト重み更新

キャッシュ状態の明示的な無効化なしに新しい重みをロードする V0 の挙動に合わせるため、チームは以下の V1 更新パスを実装しました：

await engine.pause_generation(mode="keep", clear_cache=False)
await engine_client.collective_rpc_async(
    "receive_weight_update",
    args=(request.model_dump_json(),),
)
await engine.resume_generation()

mode="keep" と clear_cache=False を使用することで、重み同期モデルが V0 参照と一致し、トレーニング後半の持続的な遅延が減少しました。

数値的パリティ：fp32 lm_head

バックエンドの修正後でも、最終的なパリティにはロジット計算に使用される数値パスを一致させる必要がありました。トレーナーは最終投影に fp32 の lm_head を使用しており、ロールアウトバックエンドもこの精度に合わせて更新する必要がありました。

この精度は重要です。なぜなら RL の更新はトークンの logprob を直接消費するため、ロジットの小さな変化がポリシー比率や KL ダイバージェンスに大きく影響するからです。この知見は MiniMax-M1 の技術報告書および ScaleRL 論文と一致しており、どちらも大規模 RL においてトレーニング／推論のトークン確率不一致を防ぐために fp32 ヘッド計算が必要な設計選択であると指摘しています。

なぜバックエンドの正確性が目的側の修正に先行するのか

トランケート重要度サンプリングや重要度比率の再重み付けといったツールは、古いまたは非同期のロールアウトを補正できますが、バックエンドを修正する前にそれらを適用すると結果が混乱します。目的側の補正が壊れた推論バックエンドの挙動を補うために使用されると、トレーニング曲線の解釈が難しくなります。

チームのアプローチにより、残る不一致は RL 目的の設計（例：非同期／オフポリシー問題）によるものであり、推論エンジンのバグではないことが保証されます。今後の改善は、ロールアウト時の明示的なビヘイビアポリシー logprob の保持や、最適化時にトレーナー側で旧ポリシー logprob を再計算するなど、非同期／オフポリシーのクリーンアップに焦点を当てます。

Sources