vime の AMD Instinct GPU 向け ROCm サポート

vLLM は、その強化学習 (RL) フレームワークである vime に ROCm サポートを統合しました。これにより、大規模な RL ポストトレーニングを AMD Instinct MI300X および MI355X GPU 上でネイティブに実行できるようになります。この統合により、ソースからのビルドの必要性を排除するプリビルドコンテナを含む、検証済みのエンドツーエンドのパイプラインが提供されます。

vime のアーキテクチャと設計

vime は、slime フレームワークに基づいた、3段階のデカップルされた訓練・推論設計を採用しています。RL プロセスを以下の 3 つの異なるコンポーネントに分離しています：

• Training (Megatron): メインの訓練ループ、パラメータ更新、およびロールアウト側への重みの同期を管理します。
• Rollout (vLLM + Router): 報酬または検証器（verifier）の信号に基づいて訓練サンプルを生成するための推論サンプリングを処理します。
• Data Buffer: 訓練とロールアウトの間のブリッジとして機能し、カスタムロールアウトロジックとプロンプトインジェクションを管理します。

RL における AMD Instinct GPU のハードウェア上の利点

RL ポストトレーニングは、訓練側の重み (Megatron 形式) と推論側の KV キャッシュ (vLLM ロールアウト) を同時に保存する必要があるため、非常にメモリ集約的です。AMD Instinct GPU はこのプロファイルに最適化されています：

• 高い HBM 容量: MI300X は 192 GB の HBM3 を提供し、MI355X は 288 GB を提供します。この大容量により、アグレッシブなテンソル並列の必要性が減少し、トポロジーが簡素化され、クラスターの利用率が向上します。
• 優れたメモリ帯域幅: MI300X は 5 TB/s 以上、MI355X は 8 TB/s の集約帯域幅を提供します。RL ロールアウトはメモリ帯域幅に制約されるため、この高いスループットは自己回帰的なトークン生成フェーズにおけるステップレイテンシを低減します。
• ネイティブなエコシステム統合: ROCm のオープンソースの性質により、vime は個別のコードパスを必要とせずに既存の vLLM ロールアウトスタックを継承でき、ROCm に対するネイティブな PyTorch および vLLM のサポートを活用できます。

ROCm での技術的実装

vime を AMD ハードウェアに導入するには、いくつかの主要なソフトウェアコンポーネントの統合が必要でした：

Megatron-LM バックエンド

vime は、ROCm と互換性のある Megatron-LM のフォークを使用しています。互換性を確保するために、非 CUDA ビルドで CUDA 融合カーネルの初期化をガードするパッチが実装されました。このパイプラインは、単一の GPU 上での HuggingFace から torch_dist チェックポイントへの変換をサポートし、勾配蓄積にはネイティブの PyTorch パスを利用します。

同一配置（Colocated）重み同期

同一配置モードでは、Megatron と vLLM は同じ GPU プールを共有します。vime は、各オプティマイザのステップ後に Megatron から vLLM へ更新された重みを同期するために、プロセス間通信 (IPC) を使用します。これは、ルーティングのために安定したプロセス一貫性のあるデバイス UUID を提供する torch.cuda.get_device_properties(i).uuid インターフェースによって ROCm 上で有効化されます。

リソース管理と Ray の統合

GPU の割り当ては HIP_VISIBLE_DEVICES を介して管理されます。vime は、CUDA_VISIBLE_DEVICES と並行してこれを設定することで、Megatron 訓練アクターと vLLM サブプロセス間の整合性を確保します。Ray の AMD GPU マネージャーは、これらのマスクを尊重するように構成されており、大規模にアクターワーカーを起動するには、ファイル記述子の制限を上げる (--ulimit nofile=1048576:1048576) 必要があります。

MI355X でのパフォーマンスベンチマーク

MI355X GPU 上での Qwen3-8B モデルを使用したテストでは、100 訓練ステップにわたって以下の結果が示されました：

• スループット: 約 4,100 tokens_per_gpu_per_second を維持。ポリシーがより予測可能な出力を学習するにつれてスループットは上昇傾向にあり、これにより vLLM がより効率的にバッチ処理できるようになりました。
• Logprob の安定性: train_rollout_logprob_abs_diff は約 0.012 で安定していました。訓練とロールアウトの対数確率の間のこの低い乖離は、重み同期が成功していることを示しており、NVIDIA ハードウェアでの結果に匹敵します。
• 報酬の収束: dapo-math-17k データセットを使用した場合、raw_reward は 0 付近から約 0.5–0.6 まで上昇し、ポリシーが検証器によって報酬を与えられる推論パターンを正常に学習したことを示しました。

サポートされている機能とロードマップ

現在のサポート

• GRPO 訓練
• 同一配置および非同期（分離された GPU プール）の訓練とロールアウト
• Megatron-LM 訓練バックエンドおよび vLLM ロールアウトバックエンド
• Qwen3 Dense および MoE モデルのサポート

将来のロードマップ

• 完全な vLLM Router および PD (Prefill-Decode) 分離サポート
• FP8 パイプラインの最適化
• AMD MoE ワークロード向けの R3 (Rollout Routing Replay)
• メモリリークと logprob の乖離に対処するための非同期訓練のパフォーマンス最適化
• マルチエージェント設定およびマルチターン・ツール呼び出しのための Agentic RL