vLLM V0からV1：RLにおける修正前の正確性

PipelineRL は、ロールアウト生成の推論エンジンとして vLLM を使用します。 推論エンジンはトークンをサンプリングし、トークンの logprobs を返します。 トレーナーはこれらの logprobs を使って、ポリシー比、KL、クリップ率、エントロピー、および報酬を計算します。 これらの logprobs の計算方法に差異があると、学習のダイナミクスが変わり得ます。 これが、vLLM V0 から V1 への移行の際に、私たちが排除する必要があった train-inference mismatch です。

TL;DR. vLLM V1は、4つの点を修正した後、vLLM V0の参照に一致しました。それらは、ロールアウトのログプロブを処理する部分、V1固有のランタイムデフォルト、インフライトの重み更新パス、そして最終投影に使うfp32 lm_head です。RL目的関数を変更する前に、バックエンドの挙動を修正しました。

参照となる実行はvLLM 0.8.5 を使用しました。V1の実行はvLLM 0.18.1 です。図1に最終結果を示します。赤い実行は最初のV1の試みであり、下で説明する修正後の緑の実行が最終的なV1の実行です。

移行の目的

vLLM V1はV0エンジンの大規模な書き直しです。したがって、移行の目標は意図的にかなり絞り込みました：

1. V1が、トレーナーが期待する形式でロールアウトのログプロブを返すことを検証する
2. 同一のワークロードをV0参照に対して再実行する
3. バックエンドの同等性が復元されてからのみ、目的関数レベルの変更を評価する

最初に目に見える症状は次に現れました：

• clamp_log_ratio_new_old_indicator
• kl_new_old
• entropy
• reward

これらのメトリクスは、GSPOのトレーニング実行から得られたもので、この実験で使用した目的関数です。この種の不一致は、PPO、GRPO、またはロールアウト側のログプロブを最適化対象の一部として扱うオンラインRLシステムのいずれでも発生し得ます。

最初のV1の実行では問題がはっきりと見えました。トレーナー側のログプロブと報酬が、トレーニングの早い段階でV0参照から離れていきました。

同じパターンはトレーナー側のメトリクスにも現れます。クリップ率は、最初の比較で最も読み取りやすいシグナルです。

失敗モード

考えられる原因を3つの層に分けました：

1. 意味の不一致：バックエンドが、トレーナーの期待するものとは異なる意味のログプロブを返す。
2. 推論パスの不一致：バックエンドが、キャッシュ、スケジューリング、またはリクエスト処理に関して異なるランタイムデフォルトを使うため、同じプロンプトでも異なる実行パスを辿る。
3. 目的関数の不一致：RL目的関数は、残存する陳腐化（staleness）やバックエンドの不一致の量に応じて修正が必要。

私たちは当初、第3のカテゴリを早すぎる段階で疑っていました。有用な診断は、最初の2つをバックエンドの挙動上の問題として扱い、まずそれらを除外したことから得られました。

V1バックエンドの修正

ログプロブの意味

最初の問題は意味の不一致でした。

vLLM V1 は、温度スケーリング、ペナルティ、top-k/top-p フィルタリングなどのログイット後処理が行われる前の、生のモデル出力からログプロブを返します。PipelineRL は、サンプラが使用するための処理済み分布からのログプロブを期待していました。

必要な設定は次のとおりでした：

• logprobs-mode=processed_logprobs

これにより、ロールアウトのログプロブに見られた明らかな平均オフセットが取り除かれました。学習曲線はそれでも、既知の良好なリファレンスに対してギャップを示していたため、次に問題があるのは推論経路でした。

政策比（policy-ratio）プロットがそれを直接示しています。V1 で processed_logprobs が有効になると、3 回すべての実行にわたって平均の政策比が 1.0 に非常に近い位置に維持されます。これにより平均バイアスの修正が確認できます。残っている不一致は、クリップ率、KL、エントロピー、そして下流の学習挙動として現れます。

Runtime Defaults

初期の V1 実行では、エンジンのバージョンと V1 のランタイムデフォルトが混在していました：

• prefix caching（プレフィックスキャッシュ）：初期実行では未設定のままだったため、vLLM 0.18.1 のデフォルトが適用されました
• async scheduling（非同期スケジューリング）：初期実行では未設定のままだったため、vLLM 0.18.1 のデフォルトが適用されました
• disable-cascade-attn に対する場当たり的な上書きがあり、起動時の kwarg のパススルーを通じて設定され、コミット済みの構成にある parity レシピの外側に存在していました

パリティ（parity）実行では、これらの選択を明示的に行いました：

vllm_config:
  use_v1: true
  vllm_kwargs:
    logprobs-mode: processed_logprobs
    enable-prefix-caching: false
    async-scheduling: false

プレフィックスキャッシュには別途の注記が値します。これは通常、固定されたモデル状態に対する、正しさを損なわない推論最適化です。しかし、このオンライン RL のセットアップでは、V0 リファレンス経路に対して、キャッシュの寿命と再利用における V1 専用の違いでした。アクタは、反復するプレフィックス、同時リクエスト、非同期スケジューリング、そして飛行中（inflight）の重み更新も同時に扱っていました。

プレフィックスキャッシュのヒットでは、キャッシュポリシーが重み更新の境界を無視している場合、更新前に計算された状態を再利用できます。プレフィックスキャッシュを無効にすることで、パリティ比較から V1 の自由度の1つが取り除かれました。

Inflight Weight Updates

重みの同期も、オンライン RL の更新モデルと一致している必要がありました。1 つの選択肢は、更新のたびにリクエストをドレインしてキャッシュをクリアすることで、V1 を V0 より厳格にすることです。これは別の疑問への答えになります。まずは、V1 が既存の V0 挙動に一致できることを確認する必要がありました。

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()

重要な点が2つあります：

• mode="keep" は wait や abort よりも、古い inflight 更新モデルにより近い動作になります
• clear_cache=False は、更新時にキャッシュ状態をそのまま残す V0 のラッパ挙動と一致します

ラグ（遅れ）は、有用なランタイム診断でした。初期の V1 経路は、修正済みの V1 実行に比べて、トレーニングの後半ほど永続的なラグをより多く抱えます。

残りのギャップ：fp32 lm_head

上記のV1バックエンドの修正により、明らかな移行（migration）の問題は解消されましたが、それでも最終的なパリティ（整合）には、ロジットを計算するのに使われる数値的な経路を一致させる必要がありました。トレーナーは最終射影（final projection）にfp32の lm_head を使用していました。ロールアウト（rollout）バックエンドは、その挙動に合わせる必要がありました。

密接に関連した問題は MiniMax-M1技術レポート にも見られます。彼らのRLの実行では、学習／推論のトークン確率が一致しない（training/inference token-probability mismatch）ことが示され、それをLM出力ヘッドに起因すると突き止め、ヘッドの計算をfp32で行うことで修正しました。

これは重要です。RLの更新はトークンのlogprobsを直接消費するためです。ロジットの小さな変更は、ポリシー比率、KL、クリッピングにおいて目に見える形で現れ得ます。そのため、最終射影の精度は、オンラインRLにおける正しさ（correctness）の表面の一部です。 ScaleRLの論文 では後に、そのRLレシピの一部としてfp32のlogits／ヘッド計算を含め、さらに大規模RLにおける有用な設計選択として、それをアブレーションしています。

fp32 lm_head の経路を含めることで、報酬は最終的なパリティ結果をコンパクトに示します。図6では、最終的なV1実行がV0の参照に追随しています。一方で、最初のV1試行では、明らかに異なる報酬曲線が得られています。

アブレーション

負の結果は重要です。よくある説明を排除できるからです。

• processed_logprobs 単独：セマンティックなlogprobのバグは修正しましたが、学習時の不一致は残りました。
• バッチ不変性：別のテストでも不一致は残り、ラグがより大きく、クリップ率も高く、そしてNCCLの合併症もありました。
• 最初のV1実行を公平なベースラインとして扱う：最初のV1実行では、V1専用のデフォルトが複数有効になっていたため、公平ではない移行（migration）の比較になっていました。

私たちがまずバックエンドの正しさを修正した理由

切り詰められた重要度サンプリング（truncated importance sampling）、重要度比の再重み付け（importance-ratio reweighting）などの、目的側（objective-side）の補正は有用なツールです。もしロールアウトが意図的に古く、非同期に生成されていたり、トレーナー側ポリシーと同値性が保証できないバックエンドによって生成されていたりするなら、何らかの補正を加えるのが正しい選択であることが多いです。

この問題で最初にあったのは推論の正しさ（inference correctness）でした。V1へ移行した後、ロールアウトバックエンドは、トレーナー側の前提を壊すlogprobsと実行時の振る舞いを返しました。その時点で目的側の補正を追加すると、2つの問いを混ぜてしまうことになっていました：

• 推論バックエンドは正しいlogprobsを生成しているのか？
• logprobsが正しいとしても、目的関数にはそれでもオフポリシー（off-policy）や非同期（async）の補正が必要なのか？

これらの問いは分ける必要があります。さもないと、目的側の補正が壊れた推論バックエンドの挙動を相殺してしまい、学習曲線が解釈しづらくなります。

ただし、現在の目的（objective）はまだ改善できます。推論のパリティが復元された後は、次の改善は通常の非同期／オフポリシーのクリーンアップです：

• ロールアウト時刻の、明示的な行動（behavior）ポリシーのlogprobsを保持する
• 最適化時に、トレーナー側の旧ポリシーのlogprobsを再計算する
• バックエンドの不一致による補正と、ポリシー更新比（policy-update ratio）を分離する
• 補正項に対するESSのような診断指標を、集約したトレーナーの指標と並行して追跡する

この移行から得られる主な教訓は、より狭いものです。まずバックエンドの正しさを直し、その後に残った不一致に対する補正を追加します。

いいね

1

•