vLLM移行でRLが不安定になる理由と4つの確認ポイント

強化学習（RL）の学習が不安定なとき、原因は報酬設計や学習アルゴリズムだけとは限りません。vLLM移行の事例では、同じモデルを使っていても、推論エンジンが返すlogprobのわずかな違いが、KLダイバージェンスやclip率、報酬などの指標を変えていました。まずはバックエンドの計算、重みの反映、数値精度を旧環境と比較することが、問題解決の出発点になります。

vLLM移行でRLの結果が変わる根本原因

vLLM移行後にRLの挙動が変わった場合、最初に確認すべきなのは推論エンジンからトレーナーへ渡るlogprobです。モデル名や重みが同じでも、計算経路や既定設定が異なれば、トレーナーが受け取る評価値は一致しないことがあります。

RLでは、モデルがトークンを生成し、その生成結果をもとに、どの行動がどれほど良かったかを計算します。logprobは、各トークンを選ぶ確率の対数値です。たとえば、モデルがある単語を選んだとき、その選択がどれほど自然だったのか、別の候補よりどれほど選ばれやすかったのかを表す材料になります。

トレーナーはlogprobを利用して、現在の方策と以前の方策の比率、KLダイバージェンス、clip率、エントロピー、報酬などを計算します。そのため、生成処理で生じた小さな差が、学習の更新量や次に生成されるデータへ連鎖する可能性があります。これは、同じ答案を採点しているつもりでも、採点基準の小さな違いが結果に影響する状況に似ています。

logprobの差が学習指標へ伝わる仕組み

KLダイバージェンスは、モデルの出力傾向が以前の状態からどれほど変化したかを見るための指標です。clip率は、更新が大きくなりすぎないように制限された割合を示し、エントロピーは出力のばらつきや不確実さを確認する材料になります。いずれもlogprobと関係するため、推論側の差を無視できません。

V0とV1の比較で確認すべき4つの論点

ServiceNowの元記事では、vLLM 0.8.5をV0の基準環境、vLLM 0.18.1をV1の環境として比較しています。最初のV1移行ではトレーナー側の指標がV0と一致しませんでしたが、4つの論点を修正し、V0のリファレンスに近づける作業が説明されています。

重要なのは、単に同じモデルを読み込めているかだけを見るのではなく、生成後のlogprobがどの形式で保存され、どの時点で処理され、学習のどの計算に使われているかを追跡することです。推論、データ受け渡し、評価、更新という一連の流れを分けて確認すると、差が発生した場所を見つけやすくなります。

移行時に見落としやすい境界

特に注意したいのは、設定ファイルに明記されていないランタイムのデフォルトです。バージョン変更によって自動選択される設定が変わると、利用者が意識していない部分で計算経路や数値の扱いが変化する可能性があります。

また、rollout（モデルに実際に文章を生成させる処理）の後にlogprobをどのように扱うかも重要です。値を保存するだけでなく、型変換、マスク、並べ替え、集約などが行われていないかを確認し、V0とV1で同じ順序と意味になっているかを確かめる必要があります。

inflight weight updateがモデルの不一致を生む理由

inflight weight updateは、推論を動かしながらモデルの重みを更新する経路です。RLでは、トレーナーが更新した重みを推論側へ反映し、その新しい重みで次のデータを生成するため、更新のタイミングと反映状態が学習の整合性に直結します。

トレーナーが新しい重みを持っているのに、推論エンジンが古い重みのまま生成を続ければ、評価を行うモデルと文章を生成するモデルが一致しません。カーナビの地図だけ更新され、車載端末には古い地図が残っている状態を想像すると分かりやすいでしょう。遅延や反映漏れが小さくても、RLではフィードバックが繰り返されるため、挙動の差として表面化します。

確認では、重み更新の開始、転送、推論側での受け入れ、次の生成への利用という段階を分けて記録します。どの更新番号の重みで生成したrolloutなのかを追跡できれば、トレーナーの想定と推論側の実体が一致しているかを判断しやすくなります。

1. トレーナーが更新した重みの識別情報を記録する。
2. 推論エンジンが受け取った更新の識別情報を確認する。
3. 更新後に生成されたrolloutとlogprobを同じ更新単位で比較する。
4. 反映前のデータが後続の学習へ混ざっていないか確認する。

fp32のlm_headがlogprobの再現性に関係する理由

lm_headは、モデル内部の表現を語彙ごとのスコアへ変換する最終出力層です。次に出す単語の候補を決める最後の投影部分にあたり、ここで計算された値が確率やlogprobのもとになります。

AIの推論では、fp16やbf16のような低精度の浮動小数点形式を使うことで、高速化やメモリ使用量の削減を目指せます。一方で、低精度化には丸め誤差が伴います。通常の文章生成だけなら見過ごせる差でも、RLのようにlogprobを何度も評価へ利用する処理では、学習指標の違いとして現れる可能性があります。

元記事の事例では、最終出力層であるlm_headをfp32（32ビット浮動小数点）で計算することが、V0の結果に近づけるための論点として扱われています。これは常にfp32が最適だという意味ではなく、旧環境との互換性や再現性を優先する場面では、どの層をどの精度で計算していたかを明示的に確認する必要があるということです。

RLの目的関数を変更する前に、まずバックエンドの正しさを確認する。

精度差を調べるときの考え方

精度差を調べる際は、生成された文章だけで判断しないことが大切です。文章が同じでも、候補トークンのスコアやlogprobが異なる場合があります。モデルの出力テキスト、トークン列、logprob、使用されたデータ型を同じ入力で記録し、比較対象を分けて見ると原因を切り分けやすくなります。

目的関数を変える前に行う移行検証の手順

V1の初期結果がV0と違っても、すぐに報酬設計、KL係数、クリッピングの仕組みを変更するべきではありません。まず同じ入力に対する推論結果、logprob、重みの反映状態、数値精度を確認し、アルゴリズムの問題とバックエンドの差を分離します。

移行検証は、一度にすべてを変更するのではなく、比較対象を固定して段階的に進めます。旧環境をリファレンスとして、同じモデル、同じ入力、同じ生成条件に近づけたうえで、出力テキストだけでなく、トークン単位の値とトレーナー側の指標を順に確認します。

1. 基準を固定する：元記事で示されたV0のvLLM 0.8.5を比較の基準として、使用モデルや入力データを記録します。
2. 生成経路を比べる：V0とV1でrolloutの形式、トークン列、logprobの保存方法を確認します。
3. 設定を明示する：V1固有のランタイムデフォルトを洗い出し、暗黙の設定を可能な範囲で明示します。
4. 重みを追跡する：更新された重みが推論側へ届き、次の生成に使われているかを確認します。
5. 出力層を確認する：lm_headの計算精度を見直し、必要に応じてfp32での比較を行います。

この手順の目的は、数値を無理に完全一致させることではありません。どの差が仕様上許容され、どの差が学習ループの意味を変えるのかを明確にすることです。比較結果をログとして残しておけば、将来のバージョンアップや別の推論バックエンドへの移行でも、同じ検証を再利用できます。

機械学習パイプライン全体への教訓

今回のvLLM移行の事例は、AI開発で不調が起きたときにモデルやアルゴリズムだけを疑わないことの重要性を示しています。データ前処理、推論設定、バージョン差、キャッシュ、重みの同期、数値精度など、土台の部分にも原因はあります。

特にRLは、推論結果を評価し、その評価を使って次の学習を進めるフィードバックループです。推論と学習の間に小さなズレがあると、そのズレを含んだデータで次の更新が行われます。だからこそ、新しいバックエンドを採用すること自体をゴールにせず、旧環境との比較テストを用意してから運用へ進む姿勢が大切です。

出典: vLLM V0からV1へ：RLでは補正より正しさを先に確認する

よくある質問