vLLM、AMD GPU で RL 学習を完全サポート

vime の発表以来、AMD チームは vime チームと緊密に連携し、ROCm への互換性サポートを提供してきました。具体的には AMD Instinct ハードウェア上でエンドツーエンドのパイプラインを検証し、ROCm 固有の不具合修正をオープンソースプロジェクトへ反映させるとともに、ビルドの手間を省ける事前構築済みコンテナも提供しています。これにより、AMD ユーザーはソースコードからのビルドなしですぐに使い始めることができます。

このたび、vLLM エコシステムの強化学習フレームワークである vime に対する ROCm サポートを、AMD Instinct MI355X GPU で正式に発表できることを嬉しく思います。このサポートにより、大規模な RL 事後トレーニングのワークフローが AMD ハードウェア上でネイティブに実行可能となり、vime の完全なパイプラインを ROCm エコシステムへ持ち込むことが実現します。

本記事では、vime とそのアーキテクチャの概要、AMD Instinct GPU が RL ワークロードに適している理由、現在の ROCm サポート状況と検証済みの機能、そして ROCm 上でエンドツーエンドの RL 学習ジョブを実行する手順について解説します。

vime の仕組み

vime は vLLM チームによって 2026 年 6 月に発表され、すぐに vLLM エコシステムにおける RL 事後トレーニングの中核として注目されるようになりました。ROCm サポートにより、これらのワークフローは AMD Instinct GPU で特別な設定なしにネイティブで実行できるようになります。

imagevime のアーキテクチャ

vime は、slime が採用している 3 つの段階に分かれた学習と推論の分離設計を踏襲していますが、ロールアウト（展開）バックエンドとして SGLang の代わりに vLLM を使用するのが特徴です。

• 学習 (Megatron): 主要な学習ループであり、パラメータの更新を担当し、重みをロールアウト側に同期します。

• Rollout (vLLM + Router): 推論サンプリングを行い、報酬や検証器の信号を用いてトレーニング用のサンプルを生成します。

• Data Buffer: トレーニング側とロールアウト側をつなぎ、プロンプトインジェクションやカスタムロールアウトロジックを管理する役割を果たします。

このパイプライン全体が、ROCm 上でエンドツーエンドで検証済みです。

なぜ AMD Instinct™ GPU を選ぶのか

強化学習（RL）のポストトレーニングは、現代の機械学習において最もメモリ集約度の高いワークロードの一つです。各トレーニングステップでは、トレーニング側の重み（Megatron フォーマット）と推論側の KV キャッシュ（vLLM のロールアウト用）を同時に保持する必要があります。コロケートモードでは、これらが同じデバイスメモリのプール内で競合します。

AMD Instinct MI300X および MI355X GPU は、以下の理由からこのワークロードプロファイルに極めて適しています。

• 大容量の統合 HBM: MI300X は GPU あたり 192 GB の HBM3 を搭載し、MI355X ではこれが 288 GB に増強されています。この余裕により、大規模モデルをトレーニング用に収容する際に、メモリ負荷を分散させるために過剰なテンソル並列化を行う必要がなくなり、トポロジーの複雑さが軽減され、クラスターの利用効率が向上します。

• 高いメモリー帯域幅: HBM3 は MI300X で合計 5 TB/s 以上の帯域幅を提供し、HBM3E を搭載した MI355X では 8 TB/s に達します。RL のロールアウト、つまり大規模な自己回帰トークン生成は、本質的にメモリー帯域幅に依存する処理です。デコードステップごとに、フル KV キャッシュとモデル重みが HBM から読み込まれます。高い帯域幅はステップのレイテンシを短縮し、RL パイプライン全体の時間を支配するロールアウトフェーズのスループットを向上させます。

• オープンソフトウェアエコシステム：ROCm は、HIP、LLM、MIOpen といったオープン標準に基づいて構築された AMD のオープンソース GPU コンピューティングプラットフォームです。vLLM と PyTorch はどちらも ROCm をネイティブでサポートしており、これにより vime も vLLM のロールアウトスタックをそのまま継承し、別の実装パスを用意する必要がありません。ROCm 環境で作業しているチームであれば、すでに慣れ親しんだツールチェーンを拡張して利用できます。

詳細なトレーニングプロセス

ROCm 上で vime を有効化するには、スタックの複数のコンポーネントを検証・統合する必要があります。AMD GPU で vime を実行した際に裏側で何が起こっているのか、以下に解説します。

• Megatron-LM トレーニングバックエンド：vime はトレーニングエンジンとして Megatron-LM を採用しています。ROCm 環境では、ROCm 対応のフォークと、非 CUDA ビルドにおける CUDA 融合カーネルの初期化を保護する小さなパッチを組み合わせて、Megatron スタックをクリーンに構築します。トレーニングループは ROCm 対応の Megatron パッチと、ROCm 固有の起動フラグを使用して実行されます。勾配の蓄積にはネイティブな PyTorch のパスが使用され、これは ROCm で完全にサポートされています。HuggingFace 形式から Megatron の torch_dist 形式へのチェックポイント変換は単一の GPU 上で行われ、スムーズに完了します。これにより生成されるレイアウトは、各トレーニング実行の開始時に Megatron が正しく読み込めるものです。

• Colocated weight synchronization: colocated モードでは、Megatron と vLLM は別々のノード分区ではなく、同じ GPU プールを共有します。オプティマイザのステップごとに、Megatron は更新された重みを IPC を経由して vLLM エンジンに同期するため、ロールアウトワーカーは常に最新のポリシーから生成を行います。これにより、ネットワーク往復が不要になります。ROCm 環境では、torch.cuda.get_device_properties(i).uuid インターフェースが安定したプロセス一貫性のデバイス UUID を返すため、vime の UUID キー付き IPC ルーティングは修正なしで正しく動作します。

• GPU visibility and Ray integration: ROCm では HIP_VISIBLE_DEVICES 環境変数を使用して GPU の割り当てを制御します。vime の起動スクリプトでは、この変数を CUDA_VISIBLE_DEVICES とともに設定しているため、Megatron トレーニングアクターと vLLM サブプロセスの両方がジョブ全体を通じて一貫したデバイス順序を見ることができます。Ray の AMD GPU マネージャーはこれらの可視性マスクを上書きしないように構成されているため、ジョブドライバーおよびすべての Ray アクター（Megatron トレーニングアクターとそのワーカープロセスを含む）が競合することなく正しい GPU セット上で動作します。また、コンテナ起動時には --ulimit nofile=1048576:1048576 を指定してファイルディスクリプタ制限を引き上げており、これは Ray が大規模なアクターワーカーを生成する際に必要とする設定です。

Getting Started: Run vime on AMD GPUs

vime は事前構築されたコンテナを用意した ROCm 対応のワークフローを提供しており、最小限の設定で完全な RL パイプラインを実行できます。

Launch the vime ROCm container

ROCm イメージの取得

docker pull vllm/vime-rocm

コンテナの起動

docker run -d --name vime --ulimit nofile=1048576:1048576 \

--ipc=host --network=host --device=/dev/kfd --device=/dev/dri \

--security-opt seccomp=unconfined --group-add video --privileged \

-e WANDB_API_KEY=$wandb_key vllm/vime-rocm

ローンチスクリプトは W&B のオンラインモードを有効にするため、有効な WANDB_API_KEY が必要です。

コンテナへの接続

docker exec -it vime bash

このコンテナには vLLM と Megatron-LM があらかじめインストールされており、vime コードベースは /root/vime に配置されています。

モデルとデータセットのダウンロード

モデル重みのダウンロード（Qwen3-8B）

hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B

学習用データセットのダウンロード（dapo-math-17k）

hf download zhuzilin/dapo-math-17k --repo-type dataset --local-dir /root/dapo-math-17k

重みの Megatron 形式への変換

Qwen3-8B のモデル設定を読み込み、変換を実行します。

cd /root/vime && source scripts/models/qwen3-8B.sh

HIP_VISIBLE_DEVICES=0 PYTHONPATH=/root/vime:/root/Megatron-LM \

torchrun --nproc-per-node=1 tools/convert_hf_to_torch_dist.py "${MODEL_ARGS[@]}" \

--no-gradient-accumulation-fusion --attention-backend flash \

--hf-checkpoint /root/Qwen3-8B --save /root/Qwen3-8B_torch_dist

注意: ROCm では、GPU の選択に CUDA_VISIBLE_DEVICES の代わりに HIP_VISIBLE_DEVICES を使用します。

RL 学習の起動

NUM_ROLLOUT=100 VISIBLE_GPUS=0,1 bash scripts/run-qwen3-8B-amd.sh

これにより、トレーニングと推論を同一環境で実行する完全な RL パイプラインが起動します。

• vLLM rollout workers
• GRPO training loop
• On-policy rollout → train → weight-sync cycle

設定に関する注意:

• VISIBLE_GPUS: 利用可能な GPU のインデックスを 2 つ指定します。スクリプトはこの 2 枚の GPU に実行を制限し、競合を防ぎます。TP=2、単一の vLLM エンジン、コロケーションモード、DP=1 を使用します。
• NUM_ROLLOUT: トレーニングステップの数です（デフォルトは 3 で、簡易テスト用）。
• 各実行には、選択した 2 枚の GPU に合わせて約 230 GB のメモリが必要です。十分な空きメモリがある GPU でのみ起動してください。

実行後の注意:

1 つの実行が完了し、異なる NUM_ROLLOUT で再実行する場合は、チェックポイントの不整合を避けるために保存ディレクトリを削除してください：rm -rf /root/Qwen3-8B_vime/

パフォーマンス結果

上記の手順書に基づき、Qwen3-4B、Qwen3-8B（Dense）、Qwen3-30B-A3B（MoE）など複数のモデルでテストを行いました。以下は、前述の Qwen3-8B 例を実行して得たパフォーマンスデータです。

MI355X 上での Qwen3-8B モデルにおけるスループットは、やや上昇傾向を示しています。

図に示す通り、MI355X 上のスループットは 100 トレーニングステップを通じて約 4,100 tokens/gpu/second を維持し、時間とともにわずかに向上する傾向が見られます。この改善は、トレーニングが進むにつれてポリシーがより予測可能な出力を生成できるようになることを反映しています。生成されるテキストの長さが短くなったり均一化されたりすることでデコードの変動が減少し、vLLM rollout エンジンがバッチ処理をより効率的に行えるようになるためです。

imageQwen3-8B モデルを MI355X で実行した際の、トレーニング側とロールアウト側の対数尤度（logprob）の絶対誤差は約 0.012 で安定しています。

この「train_rollout_logprob_abs_diff」は、学習時の確率分布と推論時の確率分布の乖離度を測る指標です。今回の実行ではこの値が 0.012 前後で推移し、わずかに低下する傾向を示しました。Megatron のトレーニングバックエンドと vLLM のロールアウトワーカー間で重み同期が行われているため、両者の整合性が保たれ、ポリシー勾配信号を壊す対数尤度のドリフト（drift）を防げています。

信頼できる GRPO 更新を行うには、この誤差が安定して低いことが必須条件です。今回のような低値は、NVIDIA ハードウェアで報告されている結果と同等の水準と言えます。

imageMI355X 上で Qwen3-8B モデルを使用した場合、ステップ 100 までに生の報酬（raw reward）はほぼゼロから 0.5〜0.6 に上昇します。この「raw_reward」はサンプリングされたトレーニングプロンプトで測定され、ステップ 0 ではほぼゼロから始まり、徐々に上昇してステップ 100 で 0.5〜0.6 に達します。

初期化直後のモデルはまだ強化学習（RL）による調整を受けていないため、数学問題に対しては事前学習時の分布のままにアプローチします。dapo-math-17k データセットの競技レベルの問題では、 freshly initialized なポリシーが解けるケースが極めて少なく、報酬もほぼゼロになります。トレーニングが進むと、モデルは部分的または完全に正解した問題から勾配シグナルを受け取り、検証器（verifier）が評価する推論パターンを好むようになります。

上昇するトレーニング報酬は、サンプリングされたプロンプト上での最適化の進捗を示しています。ただし、ROCm ランチャーでは評価（EVAL_ARGS=()）が無効化されているため、汎化能力を正確に評価するにはホールドアウトデータによる検証が必要です。

AMD における機能サポートロードマップ

現在、AMD 環境で vime のコア機能が利用可能です。具体的には以下の通りです：

• GRPO トレーニング
• トレーニングとロールアウトの共置（colocated）実行
• アクターとロールアウト用の GPU プールを分離した非同期（non-colocated）トレーニング
• Megatron-LM をバックエンドとしたトレーニング
• vLLM をバックエンドとしたロールアウト
• Qwen3 の Dense モデルおよび MoE モデルのサポート

今後、vime チームと AMD チームは、追加機能へのサポート拡大にも注力していきます。対象となる機能には以下が含まれます：

• 完全な vLLM Router および PD（Prefill/Decode）分離サポート
• FP8 パイプライン最適化

• AMD の MoE ワークロード向けに R3 (Rollout Routing Replay) を導入
• 非同期トレーニングパイプラインのパフォーマンス最適化（学習とロールアウトの対数尤度発散の改善、メモリリーク問題への対応）
• マルチターンツール呼び出しやマルチエージェント設定に対応するアジェンティック RL

当社の目標は、進化を続ける vime と vLLM のロードマップに沿った、継続的なパフォーマンスと機能の向上です。

謝辞

本成果を実現したすべての貢献者に感謝いたします：

AMD の開発者コミュニティおよび vime コミュニティの皆様

本研究を通じていただいた多大なご協力とご支援に心より御礼申し上げます。

参考文献

• vime リポジトリ: https://github.com/vllm-project/vime
• vime 発表ブログ: https://vllm.ai/blog/2026-06-09-announcing-vime
• vime AMD チュートリアル: https://github.com/vllm-project/vime/blob/main/docs/en/platform_support/amd_tutorial.md
• slime リポジトリ: https://github.com/THUDM/slime