vLLM RL-Kernel v0.1.0: ROCm 上でのビット単位のトレーニング-ロールアウト一貫性

vLLM と RL-Kernel は、AMD Instinct MI300X GPU 上でトレーニングとロールアウトの logprobs のビット単位の数値的一貫性を保証するシステムを実装しました。この統合により、トレーニングと推論を処理する異なる実行エンジンがもたらす通常の数値的ポリシーのずれが解消され、RLの後処理における重要度比や KL 診断が同一の浮動小数点計算に基づくものになります。

トレーニングとロールアウトの数値不一致の問題

強化学習（RL）ワークフローでは、トレーニングとロールアウトで通常は異なる実行エンジンと演算子が使用されます。同じモデル重みと入力を用いても、カーネル実装、並列化戦略、および還元順序の違いにより、logprobs が異なってしまうことがあります。

この不一致は、重要度比、ロールアウト-トレーニング KL 診断、およびクリッピング機構に数値的ポリシーのずれをもたらします。これまでの対策は NVIDIA プラットフォームに焦点を当てていましたが、本研究ではこれらの保証を ROCm エコシステムにも拡張しています。

技術的アーキテクチャ：vime と RL-Kernel

一貫性を達成するには、データフローのタイミングとカーネルの実際の数値実行を両方揃える必要があります。ワークロードは以下の2つのコンポーネントに分かれます：

  • vime: トレーニングのタイムラインを揃えます。どのトークンバッチがどのステップに属するかを管理し、特定のトークンを生成した重みバージョンを追跡し、新しい重みを vLLM に同期します。
  • RL-Kernel: 数値実行を揃えます。同じ値が計算に参加し、両エンジンで同じように分割・マージ・丸められることを保証します。

数値実行契約

ビット単位の一貫性を達成するため、RL-Kernel は前方パス全体に厳格な数値実行契約を課します。この契約は6つの主要な観測可能な振る舞いを定義しています：

  1. 入力依存性: 出力が依存する特定の要素と還元領域。
  2. 分割: 還元領域を部分結果に分割する方法。
  3. マージ順序: 部分結果を結合する順序。
  4. 精度: 入力、蓄積器、中間状態、出力の精度。
  5. 丸め: 丸めや下位変換が発生する正確なタイミング。
  6. プリミティブ: 使用される特定の数値プリミティブ（例：exp、log、rsqrt、SiLU、FMA）。

ROCm 上での実装

AMD Instinct MI300X にこれらの保証を適用するには、いくつかのネイティブな ROCm 実行パスを変更し、トレーニング側の挙動と一致させる必要がありました。主な技術的修正は以下の通りです：

  • GEMM カーネル: 確定的 MFMA をベースにした GEMM カーネルの実装。
  • アテンション: AITER/CK Attention の実行スケジュールの修正と、ページ化 KV レイアウトの状態問題の解決。
  • 通信: HIP IPC 通信の実装と、分散集約のランク還元順序の修正。
  • 数学とフュージョン: コンパイラのフュージョンや特定の数学関数の実装による最後のビットの違いの解消。
  • logprobs: 語彙還元の実装と、クロス-TP 語彙マージの修正。

RL-Kernel の厳格パスにおける境界修正

計算境界 厳格パスの修正
RMSNorm 還元領域、エプシロン、残差加算、出力境界
アテンション 位置、マスク、論理的ページ化 KV、分割ポリシー、LSE 精度、最終的なキャスト
GEMM と SwiGLU K 次元還元、蓄積精度、エピローグ、活性化、マテリアライズ境界
Linear logprob 実際の語彙範囲、ターゲット所有権、ローカル還元、クロスランク LSE マージ
分散集約 ペイロード所有権、dtype、固定されたランク還元順序

AMD Instinct MI300X 上での検証結果

単一ノード（8x AMD Instinct MI300X GPU）上で、Qwen3-8B モデル（BF16）を用いたエンドツーエンド実験を実施しました。トレーニングには Megatron（TP4/CP2/PP1）を、ロールアウトには2つの vLLM エンジン（それぞれ TP4）を使用しました。

主な発見

  • 不一致ゼロ: 200ステップにわたる連続したロールアウトとトレーニングステップにおいて、mismatch_count と max_abs_diff は常にゼロのままです。
  • ネイティブ vime との比較: ネイティブ vime は各ステップで不一致を示したのに対し、vime + RL-Kernel 構成はトレース全体でビット単位の同一性を維持しました。
  • 安定性: 200ステップの実行期間中、平均絶対トレーニング-ロールアウト logprob 差は常にゼロのままでした。

成果の要約

ROCm 上での vime と RL-Kernel の統合により、RLの後処理に以下の重要な機能が提供されます：

  • ビット単位の一貫性: ROCm ハードウェア上でトレーニングとロールアウトの logprobs が正確に一致。
  • 検証可能な再現性: 安定した一貫性保証により、RLトレーニングの軌道の検証が容易に。
  • 障害の局所化: 演算子のアブレーションを用いて、数値の乖離が発生する正確なシステム境界を特定可能。
  • 実行証拠: ランタイムで使用されるカーネル、HIP Graph 実行、ページ化 KV パスの包括的な検証。

Sources