AMDは、AITERを基盤とするRLスタックで、DeepSeek-V4 Flashのエンドツーエンド強化学習をAMD Instinct MI308X 24基(3ノード)上で実証しました。最初の75オプティマイザステップで、held-out AIME-2024精度は30.0%から49.2%に上がりました。
AMD InstinctはAMDのデータセンター向けAI・HPCアクセラレーター製品群で、vLLMは大規模言語モデルの推論とAPI配信のためのライブラリです。
何が新しいか
AMDは今回、推論だけや単独の学習ステップではなく、ワークフロー全体を3ノード構成のMI308Xで検証しました。対象は、長コンテキストのFP8ロールアウト、BF16 MoE actor学習、オンラインRDMA重み更新、評価、チェックポイント復旧です。学習と推論は分離したサービスとして動きます。2ノードがBF16 actorを担い、勾配計算とオプティマイザ更新の数値安定性を保ちます。専用の3ノード目がFP8 rollout engineを担い、パラメータあたりのフットプリントが小さい分、KVキャッシュと長い生成系列にGPUメモリを回します。Rayが両GPUグループを配置・調整します。
DeepSeek-V4 Flashは総284Bパラメータ、1トークンあたり約13B活性のMoEモデルで、100万トークンのコンテキストに対応します。今回の実行では、モデルコンテキストと総シーケンス長を16,384 tokensに制限しました。AMDは、この結果が報告した3ノード構成とソフトウェアバージョンに固有のものだとしています。
結果(最初の75ステップ)
| 指標 | 値 |
|---|---|
| 平均ロールアウトスループット | 311 generated tokens/s(範囲287–331) |
| オンライン報酬精度 Steps 1–25 | 45.3% |
| オンライン報酬精度 Steps 26–50 | 53.2% |
| オンライン報酬精度 Steps 51–75 | 59.2% |
| AIME-2024(Step 1、8生成/問) | 30.0% |
| AIME-2024(Step 75、8生成/問) | 49.2% |
構成
| 項目 | Actor | Rollout |
|---|---|---|
| ノード | MI308X×2(16ワーカー) | MI308X×1(8ワーカー) |
| エンジン | Megatron-LM+AMDネイティブバックエンド | vLLM AsyncLLM |
| 精度 | BF16 | FP8 per-blockの重み、FP8 E4M3 KVキャッシュ |
| 並列 | TP4/PP4/EP4、expert TP1 | TP8 |
| その他 | sequence parallelism、full activation recomputation、distributed Adam(オプティマイザ状態を全てCPUへオフロード)、BF16モーメント、256 MiB streamed-Adamチャンク | コンテキスト16,384、応答上限4,096、バッチ8,192 tokens、同時48シーケンス、GPUメモリ使用率90% |
GRPOの設定は、1ステップあたり32プロンプト×8レスポンスで、グローバルバッチは256軌跡です。学習率は1e-6、クリッピング比は0.2、上側は0.28で、明示的なKL報酬ペナルティは使っていません。rolloutとtrainで数値表現が異なるため、train/rollout比を2.0でクランプしたtoken-level importance-sampling correctionをGRPO損失に直接適用します。MoE経路では、hard-assignment R3 replayがactorの再スコアリング時にrollout時のルーティング決定を保持します。90%というメモリ目標は、オンライン方策更新用のGPUメモリを確保しつつ、長い生成系列向けのKVキャッシュ容量を残すための設定です。
再現の流れ
- 3台のMI308Xホストに各8 GPUを割り当て、Rayで2台に16 actor worker、3台目に8 rollout workerを配置します。
DAPO-Math-17Kからプロンプトをサンプリングし、vLLMで各プロンプトに8レスポンスを生成します。- レスポンスを数学報酬で評価し、Megatron actorで生成トークンを再スコアリングして、group-relative advantageを計算します。
- importance-sampling correctionを適用してBF16 actorを更新し、新しい方策をFP8へ量子化して、RCCLベースのRDMAでrollout engineへ転送します。
- AIME-2024評価を5ステップごとに実行し、分散チェックポイントを25ステップごとに保存します。保持するのは最新2つです。再開時はactorを復元してから、RDMAグループとrollout workerを再構築します。
原文は手順として、ホスト確認、ソフトウェアのバージョン固定、Dockerコンテナ起動、Rayクラスタ起動を挙げ、スモークテストの後に75ステップのジョブを実行する流れを示しています。
まだコメントはありません。