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 のチューニングは反復的に行います。通常、1 つの壁を解消すると次の壁が顕在化するため、最適なワークフローは、まずモデルを収め、次にスケールし、その後プロファイルを取り、最後に再調整することです。

最初の回答チェックリスト

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

1. モデルを収める：まず、モデルがメモリに収まるようにします。収まる範囲で最小のモデル並列化を使い、全面的な再計算より選択的な再計算を優先します。再計算と並列化でも不十分な場合にのみオフロードを追加し、大規模なレイアウトの健全性を確認するには --fake-init-process-group を使います。
2. スケールする：モデルが収まったら DP を最大化し、頻繁な通信は最速のインターコネクト内に収めます。マルチノードでのスケーリングには PP と VPP を併用し、エキスパート層では TP の追加より EP を優先します。長いコンテキストでアテンションのメモリ使用量が支配的になったら CP を追加します。
3. プロファイルを取る：支配的な壁がメモリ、通信、ホストオーバーヘッド、計算のどれかを特定します。
4. 再調整する：プロファイルで特定したボトルネックに基づいて、ディスパッチャー、オーバーラップ、FP8 モード、CUDA グラフ、再計算を変更します。
5. Parallel Folding の正確なメッシュを含めます：Attention: TP x CP x DP x PP と MoE: ETP x EP x EDP x PP。
6. デフォルトのマッピングを含めます：安全な初期起動には alltoall、H100/B200 系のシステムには flex + deepep、GB200/GB300/NVL72 系のシステムには flex + hybridep、Hopper には FP8 blockwise、Blackwell には MXFP8、またドロップレス MoE には attn、moe_router、moe_preprocess を対象とする TE スコープの CUDA グラフを使用します。

フェーズ 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 使用率が低い	グループ化 GEMM、融合処理、FP8、ディスパッチャー固有のカーネル調整

ディスパッチャーとオーバーラップの指針

ディスパッチャーの選択はボトルネックの解消に使い、最初の調整項目にはしないでください。

• moe_token_dispatcher_type="alltoall"：最も安全な初期設定で、EP サイズが小さめの場合に適しています
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="deepep"： H100 および B200 系のデプロイに適した、強力なデフォルト設定です
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="hybridep"： GB200 または GB300 NVL72 システムで最も有力な初期設定です

プロファイルで all-to-all の処理が目立つ場合は、ディスパッチャーの調整と次の設定を組み合わせてください。

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

FP8 レシピのクイック選択

プラットフォーム	推奨される開始レシピ
Hopper	FP8 ブロック単位
Blackwell	MXFP8
Blackwell、速度優先の検証	BF16 または FP8 のパスが安定した後に NVFP4

ルーターは FP32 のままにしてください。通常、大きな性能向上はエキスパートの GEMM やその他の負荷の高い行列演算から得られます。MoE の小さなコンポーネントをすべて量子化しようとすることが効果の中心ではありません。

MoE 向け CUDA グラフ

ドロップなし MoE では、部分的な TE スコープのグラフから始めてください。

• attn
• moe_router
• moe_preprocess

この方法では通常、動的なエキスパート処理をグラフの外に保ちながら、ステップ時間を大幅に短縮できます。カーネル起動のオーバーヘッドが目立つ場合は中程度の高速化が見込めますが、数 GB の追加メモリを見込み、形状が静的なまま維持されることを確認してください。

全イテレーションのグラフは、drop-and-pad のようなグラフに適したワークロードや、形状を厳密に管理した静的形状の実験にのみ使用してください。

関連資料：

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

注意点

1. 最適化の順序を誤らないでください：モデルを実行可能な状態にし、妥当な並列化を選ぶことは、マイクロ最適化より重要です。

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

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

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

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