vLLM 学習・推論整合性リリースノート
概要
LLM向けの高性能強化学習(RL)では、推論エンジンと学習エンジン間の数値的な差異により、サンプリングに使用されるポリシーとは異なるポリシーをモデルが最適化してしまう「学習・推論の不一致」がしばしば発生します。この不一致は学習を不安定にし、報酬の崩壊や再実行の増加を招く原因となります。vLLMは、ロールアウトと学習の間で**ビット単位で同一のlogprobs(差分=0)**を実現する新しいアプローチを導入しました。これにより、DeepSeek-V4-Flashにおいてベースラインのロールアウトスループットの91.3%を維持しつつ、真のオンポリシーRLを保証します。
技術的課題:学習・推論の不一致
ロールアウトはRL学習において最も時間を要する部分であり、開発者は低精度計算、低精度KVキャッシュ、および様々なカーネル最適化を採用せざるを得ません。これらの最適化は、学習時のフォワードパスと比較して、非決定論的な挙動や数値的な差異を導入することがよくあります。
DeepSeek-V4-Flashを用いた実験では、ベースラインの学習・推論間のKLダイバージェンスが100ステップで2e-3から29.9まで上昇し、報酬の低下を招きました。vLLMのアプローチは、学習時のフォワードパスにロールアウト側と同じvLLM推論カーネルを強制的に使用させることで、両エンジンが同じ重みに対して同じポリシーを計算することを保証し、この問題を解決します。
ビット単位の整合性の実現
差分を0にするために、システムは以下の4つの次元で整合性を強制します:
- カーネルの決定論性: 同じ入力に対して同じカーネルが常に同じ出力を生成することを保証し、浮動小数点加算の非結合性による問題を回避します。
- バッチの独立性: リクエストの結果がバッチ内の他のリクエストに基づいて変化する依存関係を排除します。
- フェーズの整合性: プリフィルフェーズとデコードフェーズ間の計算パスを統一します。
- GPU間整合性: 通信リダクションを修正し、タイミングやスケジューリングによって結果が変動しないようにします。
高性能カーネルをこれらの要件を満たすように修正することで、チームは低精度演算によるパフォーマンス向上を犠牲にすることなく、ビット単位で同一の結果を達成しました。
エージェント主導の実装
この整合性を新しいモデルに適用するコストを下げるため、vLLMはデバッグと調整のプロセスを自動化するエージェントスキルをリリースしました。このスキルは、エージェントを以下の手順でガイドします:
- 局所化: 2つのエンジンが乖離する最初の演算子を特定します。
- 最小限の再現: 乖離を最小限の再実行可能なテストケースに縮小します。
- 検証: 高性能なベースラインを損なうことなく、修正が整合性を維持していることを確認します。
結果とパフォーマンス
- DeepSeek-V4-Flash: 整合性バージョンは、ベースラインのロールアウトスループットの91.3%で差分0を達成しました。ベースラインの報酬が低下する一方で、整合性バージョンの報酬は安定していました。
- Nemotron-3.5 (Mambaハイブリッド): 整合性バージョンはSWE-bench Proスコア0.6を維持しましたが、ベースラインのスコアは0.56から0.22に低下しました。これは、整合性が複雑なアーキテクチャにおける学習の崩壊を防ぐことを示しています。
オープンソースでの提供
プロセス、レシピ、およびコードはコミュニティに公開されています。これには、自動調整のための train-inference-consistency スキル、DeepSeek-V4のレシピ、および Nemotron-3.5-Lightning-30B-A3B の実装が含まれます。現在、DeepSeek-V4.1の開発が進められており、すでに差分0を達成し、パフォーマンス最適化の最終調整段階にあります。