DeepSeek-V4 Flash強化学習、AMD MI355Xに対応

著者：AMD & Miles Team

AMDとMilesチームは、DeepSeek-V4 Flash RLがROCm™搭載のAMD Instinct™ MI355X GPU上でMilesサポートを受けられるようになったと発表した。これにより、DeepSeek-V4 FlashはMilesフレームワーク内でエンドツーエンドの強化学習トレーニングを実施できる。具体的には、SGLangがrollout生成とlog probabilitiesの記録を担当し、Megatronがactorトレーニングと方策更新を担当し、Milesが更新後の重みをオンラインrolloutエンジンへ継続的に送り返す構成となっている。

この取り組みは単に「動かす」だけにとどまらない。強化学習では、rollout側とトレーニング側が同一の方策を十分に一致した形で実装している必要がある。そうでなければ、同一トークンの確率推定がずれ、GRPOなどのトレーニングプロセスの安定性に影響する。DeepSeek-V4 Flashはhybrid compressed attention、mHC residual mixing、MoEルーティングなどのアーキテクチャを導入しており、SGLangとMegatron間の動作整合がより困難になっている。

主な結論

• DeepSeek-V4 Flash RLがMiles + AMD Instinct MI355X上で動作可能になった。チームはROCm環境でのエンドツーエンド実行に必要なモデル整合とオンライン更新の問題を解決した。
• 4ノードでの検証が完了。100ステップを超えるoptimizerの実行において、train-rollout間のlog-probability差異は有界に保たれ、online rewardが改善し、オフラインのAIME-2024評価スコアも向上した。
• 次フェーズの焦点は性能最適化。今後の作業には、低精度トレーニング、エンドツーエンドの性能チューニング、より大規模なクラスターへのスケールアップが含まれる。

DeepSeek-V4 Flashアーキテクチャ概要

DeepSeek-V4 Flashは2840億パラメータのMoEモデルであり、各トークンで約130億パラメータが活性化される。今回使用した構成は43層のdecoder、256個のrouted expertsとtop-6 selectionを採用し、4系統のmHC residual streamsを備える。アテンション機構については、128トークンのsliding windowと圧縮された長文脈アテンションを組み合わせている。

具体的には、C4 layersは4:1に圧縮されたKVシーケンスからtop 512エントリを選択し、C128 layersは128:1圧縮シーケンス上で稠密なattentionを実施する。mHCやMoEルーティングと合わせ、これらはSGLangとMegatronが一致した実装を維持しなければならないアーキテクチャ上の重要なパスである。

図1：DeepSeek-V4の構造概略図。mHC residual mixing、hybrid compressed attention、top-6 MoEルーティングを示している。

Miles強化学習スタックの仕組み

Milesは非同期トレーニングループを調整する。SGLangが候補回答とrollout log probabilitiesを生成し、Megatronが同一バッチのシーケンスを再スコアリングしてpolicy updateを計算しactorをトレーニングする。その後、MilesはMegatronで更新された重みを稼働中のSGLang workersに送り返し、次のrolloutに使用させる。

現在のFP8パスでは、rolloutはFP8 Hugging Face checkpointを使用し、actorトレーニングはBF16 Megatron torch_dist checkpointを使用する。つまり、2つのエンジンが異なる実行フォーマットで同一の方策を表現するため、checkpointの変換、再スコアリング、オンライン重み更新はいずれも正確性の境界に含まれる。

図2：PromptがSGLang rolloutに入り、軌跡とlog probabilitiesがMilesを経由してMegatron actorへ流れ、更新された重みがSGLangに返される。

課題1：train-rollout間のlog-probability差異の縮小

RLトレーニングは、SGLangとMegatronが同一バッチの生成トークンに対して近似した確率を算出するという前提に依存している。そのためチームはtoken-identical comparison workflowを構築した。SGLangが一度シーケンスを生成し、同一のトークンに対して2つのエンジンがスコアリングを行い、高コストな多ノード検証の前にトークンレベルの差異からモデル実装の不一致を特定する。

このフローにより、DeepSeek-V4固有の2つの差異が発覚した。初期のhash-routed MoEパスと、mHC residual mixingである。チームはMegatronのhash-routingの動作をSGLangに合わせ、Megatron側のmHC post-mixロジックを修正することで、2つのエンジンが同一のモデル意味論を保持するようにした。これらの修正によりrolloutとtrainingの数値一致性が大幅に向上した。

課題2：オンライン更新における量子化の意味論の保持

RLトレーニングでは、rollout serverは再起動なしに繰り返し更新されたpolicy weightsを受け取る必要がある。量子化モデルにおいては、重みの転送成功が更新の正確性を意味しない。packed weights、scale tensors、量子化状態に依存するランタイム情報はいずれも、更新後も正しい意味を保持し続けなければならない。

FP4およびE8M0テンソルについては、AMDが更新パスをdatatype-aware対応にし、更新後にテンソルが誤解釈されて無効な生成が発生するのを防いだ。FP8については、Miles自体が更新後のライフサイクルを定義しており、AMDはROCmスタック内でSGLangの不足しているインターフェースを補完し、rollout再開前にMilesが必要な量子化処理を実行できるようにした。

この部分での核心的な教訓は、オンライン更新は単なるバイトのコピーではなく、モデルの量子化状態を復元しなければならないということだ。

課題3：ROCm上での安定した多ノード並列戦略の確立

DeepSeek-V4 Flash RLを複数のAMD Instinct MI355Xノードに拡張する際、チームは相互に関連する2つの問題に直面した。1つは2840億パラメータのMoEモデルに対して4Kコンテキストで実行可能なmodel-parallel戦略をどう選ぶか、もう1つは多ノードのcollective communicationを安定させる方法である。

tensor parallelismを高くすると単一GPUのメモリ負荷は下がるが、collectiveのトラフィックが増加する。初期の一部多ノード構成では、tensor-parallel all-reduceやexpert all-to-allなどのRCCL collectivesが完了せずにスタックし、communication watchdogに捕捉されて実行が中断される問題が発生した。

チームは最終的に、メモリ実現可能性と通信安定性を両立するレイアウトに収束した。4つの8カードノード上でtensor-parallel 1 / pipeline-parallel 4 / expert-parallel 4を採用し、activation recomputation、optimizer stateのホストメモリへのオフロード、制御された単一GPUのトークンバジェットを組み合わせた。tensor-parallel all-reduceの負荷を減らしてpipelineとexpert parallelismに転換し、チューニングされたRCCL transport settingsと組み合わせることで、100ステップを超えるoptimizerの実行においてトレーニングをエンドツーエンドで安定させることに成功した。

4ノードAMD Instinct MI355X検証結果

チームは4つの8カードAMD Instinct MI355XノードでFP8パスを検証した。そのうち2ノードをSGLang rollout用、2ノードをMegatron actorトレーニング用とした。MilesがGRPOスタイルのトレーニングに類似したフローを調整し、トレーニングタスクは長文脈数学データセットDAPO-Math-17K、コンテキスト長は4Kとした。モデル並列構成はtensor-parallel 1 / pipeline-parallel 4 / expert-parallel 4。10ステップごとにAIME-2024でオフライン評価を実施し、各問題で8サンプルを採取した。rolloutモデルはFP8を使用し、actorはBF16でトレーニングした。

正確性とonline reward

正確性の主要指標の1つは、rolloutとtrainingが同一バッチの生成トークンに対して近似した確率を算出するかどうかである。記録されたトレーニングステップにおいて、平均絶対log-probability差異は約0.09であった。図3から確認できるように、100ステップを超え、複数回のオンライン重み更新を経ても、この差異は有界に保たれ、継続的な上昇も更新後の急激な拡大も見られなかった。チームはこれを最終的な閾値ではなく、有望なbring-up結果として強調している。

図3：最初の100トレーニングステップにおけるtrainとrollout間の絶対log-probability差異。

log-probabilityの一致に加え、online rewardもトレーニングとともに改善した。ある拡張実行において、online raw rewardは横ばいではなく明確な上昇傾向を示した。実行の前半1/3と後半1/3を比較すると平均値が向上しており、actorが継続的なGRPOトレーニングと繰り返しのオンライン重み更新によって確実に改善されており、単に既存の報酬水準を維持しているだけではないことが示された。

図4：100トレーニングステップにおけるonline raw reward。ステップごとの値、移動平均、線形フィットを含み、全体の傾きは正である。

AIME-2024オフライン評価

オンラインpass rateはトレーニング負荷下で測定され動的サンプリングの影響を受けるため、チームは独立したオフラインベンチマークも導入した。10ステップごとにAIME-2024で評価を実施し、各問題で8サンプルを採取した。これはより信頼性の高いモデル品質指標とみなされている。

最初の100トレーニングステップにおいて、オフラインAIMEのpass@1は0.39から0.49に向上し、pass@8は0.53から0.67に向上した。また、4,096トークン上限の影響で打ち切られた応答の割合は60%から55%に低下した。単一回答の正確性と多サンプルカバレッジが同時に向上したことは、GRPOがもたらしたのが単純な分布の「鋭化」ではなく、より真の能力向上であることを示している。AIME-2024は問題数が30問と少ないため単点評価にはノイズが生じやすく、全体的な傾向に注目すべきである。

図5：最初の100 RLトレーニングステップにおけるAIME-2024 pass@1/2/4/8。10ステップごとに評価し、各問題で8サンプルを採取。

次のステップ：動作可能から高性能へ

今回の検証の焦点は、DeepSeek-V4 Flash RLをAMD Instinct MI355XとROCmエコシステム上でエンドツーエンドに動作させ、重要な正確性パスが安定して機能することを実証することにあった。次のフェーズでは、低精度トレーニングのさらなる探索、エンドツーエンドのスループットとレイテンシの最適化、より大規模なクラスターへのスケールアップなど、性能最適化に重点を移していく。

総括すると、この取り組みはMiles、SGLang、MegatronがAMD GPUプラットフォーム上で連携して大規模MoEの強化学習トレーニングを実行できることを示した。ROCmエコシステムで超大規模モデルのトレーニングとアライメントを行おうとするチームにとって、DeepSeek-V4 Flashのbring-up経験は重要な参考事例となる。モデル意味論の整合、量子化状態の復元、通信並列戦略のいずれも欠かせない要素である。