nemo-mbridge-perf-moe-optimization-workflow

MoE トレーニング最適化ワークフロー

安定版ドキュメント：@docs/training/moe-optimization.md カード：@skills/nemo-mbridge-perf-moe-optimization-workflow/card.yaml 出典：Megatron Core による MoE モデルのスケーラブルなトレーニング

クイックリファレンス

まず、論文で示されている3つの壁を確認します。

• メモリの壁
• 通信の壁
• 計算効率の壁

実運用での診断では、計算効率の壁を計算とホスト／起動のボトルネックに分けてください。それぞれ必要な根拠と対処法が異なります。MoE のチューニングは反復的に行うため、次の順序で進めます。

freeze the measurement contract -> fit -> scale -> profile -> retune -> validate

最初に確認するチェックリスト

MoE 最適化ワークフローに関するプロンプトには、次の順序で回答します。

1. 測定条件を固定する：モデルとタスク、ハードウェアとトポロジー、コンテナーとコミット、データとルーティングのセマンティクス、精度、シーケンスとバッチの形状、並列化、グラフのスコープ、定常状態のメトリクス測定区間を正確に記録します。各候補には、トレーニング同等かベンチマーク専用かを明記します。
2. 適合させる：まずモデルをメモリに収まるようにします。収まる範囲で最小のモデル並列化を使用し、全面的な再計算よりも選択的再計算を優先します。再計算と並列化でも不十分な場合にのみオフロードを追加し、大規模な構成の妥当性確認には --fake-init-process-group を使用します。
3. スケールする：モデルが収まったら DP を最大化し、頻繁に使う通信は最速のインターコネクト内に収めます。複数ノードへのスケーリングには PP と VPP を使用し、エキスパート層では TP の追加より EP を優先します。長いコンテキストによってアテンションのメモリ使用量が支配的になる場合は CP を追加します。
4. プロファイルする：支配的な壁、すなわちメモリ、通信、ホストオーバーヘッド、計算のいずれかを特定します。
5. 再調整する：プロファイルで特定したボトルネックに基づき、一度に1つの変数を変更します。ディスパッチャー、オーバーラップ、低精度、CUDA グラフ、再計算は候補であり、ハードウェアに応じて自動的に選ぶ設定ではありません。
6. 検証する：候補を絞り込むには、条件をそろえた短時間の評価を行い、勝者の設定では少なくとも50ステップ実行します。要求したバックエンドまたはグラフリプレイが実際に動作したことを確認し、宣言したウォームアップ後の測定区間を計測して、損失の状態、スキップ／NaN 反復、メモリ、ステップ時間、モデル TFLOPS/GPU を報告します。
7. Parallel Folding の正確なメッシュを含めます：Attention: TP x CP x DP x PP および MoE: ETP x EP x EDP x PP。
8. 安全な初期立ち上げには alltoall を使用します。その後、パッケージと対象トポロジーが対応している場合に、flex + deepep と flex + hybridep を A/B テストします。まず BF16 とイーガー実行から始め、プロファイリングによる根拠が得られた場合にのみ、低精度または必要最小限の CUDA グラフスコープを導入します。

フェーズ0：測定条件を固定する

単一の検証対象変数を除き、次の条件を固定しない限り、比較は有効ではありません。

• Bridge、MCore、Transformer Engine、コンテナー、CUDA、NCCL のバージョン
• GPU の台数と SKU、ノードのトポロジー、ランチャー／環境設定
• モデル、タスク、データパス、シーケンス長、MBS、GBS、オプティマイザー設定
• TP、PP、VPP、CP、EP、ETP、DP のレイアウト
• ルーティングのセマンティクス、精度、再計算、ディスパッチャー、オーバーラップ、グラフのスコープ
• ウォームアップと定常状態の測定区間

次の2つの受け入れクラスを分けます。

• トレーニング同等の変更では、意図したルーティング、損失、データ、オプティマイザー、チェックポイント／再開の動作が維持されます。
• 強制的な負荷分散などのベンチマーク専用の変更は、カーネルを制御された条件で調査するのに役立ちますが、本番トレーニングの正しさや収束を裏付けることはできません。

フェーズ1：実行をメモリに収まるようにする

スループットを追求する前に、確実に収まる構成から始めます。

推奨する順序：

1. 収まる範囲で、モデル並列化を最小限にします。
2. 全面的な再計算に切り替える前に、選択的再計算を有効にします。
3. 再計算と並列化だけでは不十分な場合にのみ、オフロードを追加します。
4. クラスターの時間を無駄にする前に、単一 GPU 上で大規模な並列レイアウトの妥当性を確認するために --fake-init-process-group を使用します。

再計算の指針

MoE の実行では、選択的な再計算を優先してください。

• 最初に試す候補：layernorm、core_attn、moe_act、mlp、または モデル固有のモジュール（shared_experts、mla_up_proj）
• 実行がまだメモリに収まらない場合にのみ、全体を再計算してください
• CUDA グラフを有効にした後、再計算の設定を見直してください。一部のグラフスコープと 全体の再計算パスは相性がよくないためです

経験則として、きめ細かな再計算では必要なメモリの大部分を確保できることが多く、 全層再計算よりもスループットを再計算なしのベースラインに近く保てます。

フェーズ 2：スケールに合わせて並列化方式を選択する

優先順位：

1. モデルが収まったら、DP を最大化します。
2. 可能な場合は、頻繁に使われる通信経路を高速インターコネクト内に保ちます。
3. マルチノードでのスケーリングには PP を使い、必要に応じて VPP も使います。
4. エキスパート層では、追加の TP より EP を優先します。
5. シーケンス長によってアテンションのメモリ使用量が支配的になったら、長いコンテキストに CP を追加します。

Parallel Folding

Parallel Folding はアテンションと MoE の並列化を分離するため、単一の妥協案となるレイアウトを選ぶ必要がありません。

Attention: TP × CP × DP × PP
MoE:       ETP × EP × EDP × PP

主な設定項目：

• --expert-model-parallel-size
• --expert-tensor-parallel-size

アテンションではある程度の TP または CP が望ましい一方、エキスパート層では密結合層が許容できる範囲を超えて EP の次数を大きくするメリットがある場合に使用してください。

フェーズ 3：支配的なボトルネックをプロファイルする

ボトルネック	発生状況	主な対処法
メモリ	大量の全体再計算を行わないと実行が収まらない、またはウォームアップ中に OOM が発生する	選択的再計算、FP8、オフロード、より適切な PP レイアウト
通信	Nsight で大きな all-to-all または collective ブロックが確認される	DeepEP または HybridEP、EP のオーバーラップ、DP/TP のオーバーラップ、より適切な PP レイアウト
ホストのオーバーヘッド	GPU の空き時間、カーネル起動がボトルネックとなるトレース、Python のオーバーヘッド	CUDA グラフ、--manual-gc、より大きな MBS、CPU アフィニティの調整
計算	通信とホストの問題に対処した後も SM 使用率が低い	grouped GEMM、フュージョンの改善、FP8、ディスパッチャー固有のカーネル調整

誤読を避けてオーバーラップをプロファイルする

評価指標にはプロファイルなしの定常イテレーションを使い、原因の分析には条件を合わせたプロファイルを使います。

1. オーバーラップまたはディスパッチャーの変数を一度に 1 つだけ変更します。ルーティング、グラフスコープ、 並列化方式、バッチ形状、ランタイムは固定します。
2. 通信カーネルと計算カーネルの区間の和集合を作成し、その共通部分を測定して、隠蔽された通信量を定量化します。
3. カーネルの実行時間を合計してウォールタイムと見なさないでください。SM または帯域幅の競合により、 カーネルの実行時間が長くなる場合がありますが、GPU が実際に稼働している区間の和集合とエンドツーエンドのステップ時間は短縮することがあります。
4. トレースを、dispatch/combine の NVTX 範囲、定常状態のステップ時間、モデル TFLOPS/GPU、損失の有限性、スキップ数／NaN 数、ピークメモリと照合します。

条件を制御した 16×H100 Qwen3 30B-A3B HybridEP の実行では、通常の EP オーバーラップにより、 GEMM／アテンションで隠蔽される通信量が 0.11% から 36.55% に増加しました。プロファイルなしのステップ時間は 24.7138s から 20.9920s に短縮され、スループットは 244.039 から 287.305 model TFLOPS/GPU に向上しました。delay_wgrad_compute は無効のままです。

フェーズ 4：根拠に基づいて再調整する

プロファイルで特定したボトルネックに対処する候補のうち、最小限のものを選び、変数を一度に 1 つだけ変更します。

ディスパッチャーとオーバーラップに関する指針

ディスパッチャーの選択は、ハードウェアの対応表としてではなく、ボトルネックを解消する手段として使ってください。

• moe_token_dispatcher_type="alltoall"：最も安全に導入できる方法で、小規模な EP サイズに適しています
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="deepep"： DeepEP がインストールされ、通信が利用可能な場合の候補です
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="hybridep"： HybridEP がインストールされている場合、NVL8 と NVL72 の両システムで検討できる、トポロジー依存の候補です

HybridEP と通常の EP オーバーラップの組み合わせは、標準的な 16×H100 Qwen3 30B-A3B の構成で現在測定されている最速の方法です。一方、標準的な 256×H100 Qwen3 235B のレシピでは、標準の alltoall とオーバーラップを使用します。対象のコンテナでバックエンドの互換性とスループットをベンチマークしてください。GPU の名称や EP の次数だけでは、最速の方法は決まりません。

プロファイルで All-to-All 経路が確認できる場合は、ディスパッチャーの調整と次の設定を組み合わせてください。

• --overlap-moe-expert-parallel-comm
• --overlap-grad-reduce
• --tp-comm-overlap

まずは、通常の EP オーバーラップ、共有エキスパートのオーバーラップ、遅延させた重み勾配の計算を、それぞれ別個の候補としてテストしてください。あるモデルで効果のあった要素でも、組み合わせると性能が低下する場合があります。

低精度候補のマトリックス

検証済みの BF16 ベースラインから始めてください。ハードウェアの能力によって決まるのは、どの低精度候補を使用できるかだけであり、高速化が保証されるわけではありません。

プラットフォーム	BF16 が安定した後の候補
Hopper	対象スタックでサポートされるテンソル単位、動的スケーリング、またはブロック単位の FP8
Blackwell	MXFP8 またはその他のサポート対象の FP8 レシピ
Blackwell、速度優先の検討	BF16／FP8 経路が安定した後の NVFP4

ルーターは FP32 のままにしてください。最大の性能向上は通常、エキスパート GEMM やその他の負荷の大きい行列演算から得られます。MoE の小さなコンポーネントをすべて量子化しようとする必要はありません。意図したカーネルが実行されたことを示すログまたはトレースを必須とし、理論上のピーク FLOPS ではなく、エンドツーエンドの定常ステップ時間で候補を評価してください。

MoE の CUDA Graphs

プロファイルでホスト処理／カーネル起動の間隙が十分に大きいと確認できた後にのみ、CUDA グラフを使用してください。トークンをドロップしない MoE では、対象範囲を TE に限定した最小の部分グラフ候補から始めてください。

• moe_router
• moe_preprocess

モデルでサポートされ、同じ条件を揃えたスタックで性能が向上する場合に限り、attn を追加してください。キャプチャに成功しても高速化の証拠にはなりません。また、ディスパッチャー、オーバーラップ、または精度を変更すると、グラフによる性能向上が失われることがあります。

この方法では、動的なエキスパート処理はグラフの外に保たれます。追加メモリを見込み、シェイプが静的なままであることを確認してください。キャプチャだけでなく再生も確認し、キャプチャ後のイテレーションのみを計測してください。

全イテレーションのグラフは、トークンのドロップとパディングを行うワークロードや、静的シェイプを厳密に管理する実験など、グラフに適したワークロードでのみ使用してください。

関連資料：

• @skills/nemo-mbridge-perf-cuda-graphs/SKILL.md
• @docs/training/cuda-graphs.md
• @docs/training/activation-recomputation.md

フェーズ 5：検証とエビデンスのまとめ

ワークロードが許す場合は、手軽なスクリーニングにウォームアップ後の 6〜12 イテレーションを使用してください。選定した候補では、少なくとも 50 ステップを実行し、ステップ 41〜50 など、固定した定常区間を報告してください。最終的なエビデンス一式には、次の内容を含めてください。

• 正確なコマンド／設定の差分、コミット、コンテナ、ハードウェア、トポロジー
• ルーティング／データの意味論の明示、およびトレーニング相当かベンチマーク専用かのラベル
• 意図したディスパッチャー、精度カーネル、オーバーラップ、グラフ再生が有効だったことを示す証拠
• 同じプロファイルなしの定常区間におけるステップ時間と、GPU あたりのモデル TFLOPS
• 有限値の損失、スキップ数／NaN 数、ピークメモリ、および本番経路で必要な場合はチェックポイント／オプティマイザー状態の検証結果
• 主張する因果メカニズムに対応した、条件を揃えた A/B プロファイルの証拠

複数の変更を加えた最終候補の合計性能向上を、以前の単一の A/B テストに帰属させないでください。たとえば、Qwen3 のオーバーラップ実験では、GPU あたりの TFLOPS が 244.039 から 287.305 に増加することを個別に確認しました。その後、HybridEP を追加調整した標準レシピでは 299.352 に達しました。これらは別の問いに答える結果です。

落とし穴

1. 最適化する順序を間違えないでください：モデルの適合と適切な並列度の選択は、細かな最適化よりも重要です。

2. プラットフォームによってボトルネックは変わります：H100 クラスの実行では通信がボトルネックになりやすい一方、GB200 や GB300 では CPU や起動オーバーヘッドが早い段階で表面化することがよくあります。

3. FP8 の MFU は実際より低く見えることがあります：精度モードを切り替えるときは、MFU だけでなく絶対スループットも比較してください。

4. CUDA グラフと再計算は相互に影響します：TE スコープのグラフは通常、全面的な再計算ではなく、選択的な再計算と組み合わせます。

5. 大規模では Parallel Folding は必須です：アテンション層とエキスパート層で明らかに異なるレイアウトが必要になると、単一の共有 TP または EP プランは双方に負担をかけます。

6. カーネル時間の合計は、表面化している時間ではありません：オーバーラップを検証するときは、区間の和集合と通信・計算の交差を使用してください。

7. ベンチマーク専用のセマンティクスは、本番環境での合格基準ではありません：ルーティングの強制、合成データの使用、オプティマイザーやチェックポイント処理経路の無効化は明示し、トレーニングと同等の結果とは別に検証してください。

8. 機能の有効化には根拠が必要です：バックエンドがフォールバックする場合、グラフをキャプチャできても効果がない場合、または低精度のレシピで意図したカーネルを使用できない場合は、設定ダンプだけでは不十分です。

最新の署名更新日：2026-08-03。