nemo-mbridge-perf-moe-comm-overlap

MoE 通信のオーバーラップ

概要については、以下を参照してください。

• @docs/training/communication-overlap.md
• @skills/nemo-mbridge-perf-moe-comm-overlap/card.yaml

すぐに判断するには

次の場合は、MoE 通信のオーバーラップを使用します。

• EP > 1
• プロファイルでトークンのディスパッチまたは結合にかかる時間が確認できる
• 実行がすでに正しく動作しており、スループットを調整する段階である

初期の立ち上げ手順として有効にするのは避けてください。ディスパッチャー、ルーティングモード、再計算プランが安定してから検証するほうが簡単です。

有効化

cfg.comm_overlap.overlap_moe_expert_parallel_comm = True

# Optional: delayed wgrad for additional overlap
cfg.comm_overlap.delay_wgrad_compute = True

# IMPORTANT: disable shared expert overlap when using dispatch overlap
cfg.model.moe_shared_expert_overlap = False

前提条件

• expert_model_parallel_size > 1
• num_moe_experts > 1
• moe_token_dispatcher_type は "alltoall" または "flex" である必要があります
• 精度：BF16 または FP16
• PP を使用する場合、VPP（virtual_pipeline_model_parallel_size）を設定する必要があります（None 以外）

Flex ディスパッチャーの有効化

moe_flex_dispatcher_backend を設定するだけでは、flex ディスパッチは有効になりません。 moe_token_dispatcher_type = "flex" も設定する必要があります。

再計算と CUDA グラフの相互作用

• 完全再計算は、オーバーラップ経路との組み合わせには適していません。
• CUDA グラフのスコープにアテンションまたは MoE ルーターの処理が含まれる場合、delay_wgrad_compute によってさらに制約が生じます。
• 実際には、オーバーラップを有効にする場合は選択的再計算との組み合わせのほうが安全です。

測定結果

HybridEP の実運用規模での検証

2026-07-25 に実施した、Qwen3 30B-A3B の事前学習に関する制御比較では、16 基の H100 GPU、BF16、シーケンス長 4096、TP=1、PP=1、CP=1、EP=16、MBS=1、GBS=1024、強制バランスルーティング、HybridEP、および Transformer Engine の CUDA グラフスコープ moe_router と moe_preprocess を使用しました。性能に関する唯一の変更は通常の EP オーバーラップであり、勾配重み計算の遅延は無効のままでした。

ケース	定常状態の区間	ステップ時間	モデル TFLOPS/GPU
EP オーバーラップなし	反復 5-20	24.7138s	244.039
EP オーバーラップあり、探索実行	反復 5-20	21.0725s	286.208
EP オーバーラップあり、独立検証	反復 41-50	20.9920s	287.305

独立検証では、再現したベースラインに対してステップ時間が 15.059% 短縮し、スループットが 17.729% 向上しました。損失は有限値を維持し、反復のスキップや NaN はなく、rank-0 のピーク割り当てメモリは 62.166 GiB でした。

同一手法で取得した rank-0 の Nsight Systems 比較では、それぞれのケースで 463,348 個のカーネルを記録しました。

プロファイル指標	オーバーラップなし	オーバーラップあり
GEMM／アテンションと並行する通信	9.079ms	3,958.997ms
計算によって隠された通信時間	0.11%	36.55%
GPU アクティブ区間の合計	22.821s	21.221s
HybridEP のディスパッチ＋並べ替え NVTX	4.253s	1.767s
HybridEP のメタデータ前処理 NVTX	3.109s	0.670s

これは、性能向上が、ディスパッチャー、ルーティング、グラフスコープ、バッチ形状、並列レイアウトを変更したことではなく、露出していた HybridEP のディスパッチ／結合処理を隠したことによる、直接的な証拠です。

正しさを優先した alltoall のスモークテスト

2026-05-18 に current-main で実施した、Qwen3 30B-A3B のモック事前学習に関する H100 x16 のスモークテストでは、EP=16、alltoall、グローバルバッチサイズ 1024、CUDA グラフ無効を使用しました。以前の立ち上げで PyTorch 25.11／TE／Triton のスタックが Transformer Engine の融合パーミュテーションで失敗したため、moe_permute_fusion=false としました。

結果はリリース品質の検証ではなく、傾向を示すものでした。

• EP オーバーラップなし：反復 3-8 の定常状態の平均は 41.25s
• EP オーバーラップあり：反復 3-8 の定常状態の平均は 31.31s
• EP オーバーラップと delay_wgrad_compute の併用：反復 3-8 の定常状態の平均は 31.20s

通信が露出している場合、EP オーバーラップがノード間の alltoall MoE の形状で有効な可能性を示す証拠として扱ってください。勾配重み計算の遅延が独立した改善をもたらす証拠ではなく、融合パーミュテーション経路を検証するものでもありません。同じ形状で以前に実施した 2026-05-16 の短時間スモークテストでも、同じ傾向が確認されました。

コードの参照先

• オーバーラップの検証：src/megatron/bridge/training/comm_overlap.py
• Flex ディスパッチャーのバックエンド：src/megatron/bridge/training/flex_dispatcher_backend.py
• 設定：src/megatron/bridge/training/config.py
• 単体テスト：tests/unit_tests/training/test_comm_overlap.py
• DeepEP テスト：tests/unit_tests/training/test_deepep.py

注意点

1. 共有エキスパートのオーバーラップの競合：moe_shared_expert_overlap と overlap_moe_expert_parallel_comm は競合する場合があります。ディスパッチの オーバーラップ経路を使用する場合は、共有エキスパートのオーバーラップを無効にしてください。

2. VPP なしの PP：パイプライン並列が有効な場合、MoE のオーバーラップには VPP が必要です。 VPP がないと、オーバーラップのスケジューリングで正しくインターリーブできません。

3. Flex とバックエンドフラグは別物：moe_token_dispatcher_type が "alltoall" のままでは、 moe_flex_dispatcher_backend="deepep" だけを設定しても何も起こりません。

4. 保守的なレシピのデフォルト設定：ほとんどの公開レシピでは MoE のオーバーラップは 無効になっています。オーバーライドで明示的に有効にする必要があります。

5. 性能向上はワークロードによって異なります：オーバーラップは、ディスパッチ通信が ステップ時間の中で明確な割合を占めている場合に最も効果があります。小規模または負荷の 低い EP 実行で効果があるとは限りません。

6. カーネル時間の合計は実時間ではありません：カーネルが同時実行されると、SM や帯域幅を 取り合うために実行時間が長くなることがあります。そのため、オーバーラップによって、 ストリームごとのカーネル時間の合計が増える一方で、露出する時間区間の和集合と エンドツーエンドのステップ時間が短くなることがあります。

検証

初期化中に、オーバーラップ関連のログメッセージを確認してください。comm_overlap.py の 通信オーバーラップ検証では、前提条件が満たされていない場合に例外が発生するため、 エラーなく起動すれば、この機能が有効であることを確認できます。

短時間の性能ハーネスのスモークテストでは、コマンドの形式を明示し、 一度に変更するオーバーラップ設定は 1 つだけにしてください。

uv run python scripts/performance/run_script.py \
  -m qwen \
  -mr qwen3_30b_a3b \
  --task pretrain \
  -g h100 \
  -c bf16 \
  -ng 16 \
  -gn 8 \
  --max_steps 8 \
  --cuda_graph_impl none \
  --moe_flex_dispatcher_backend None \
  --moe_a2a_overlap false \
  --tokenizer_type NullTokenizer \
  comm_overlap.overlap_moe_expert_parallel_comm=true \
  comm_overlap.delay_wgrad_compute=false \
  model.moe_shared_expert_overlap=false

初期検証中に MoE の融合パーミュテーションで失敗する場合は、 model.moe_permute_fusion=false を追加して、オーバーラップの計測とランタイムスタックの 検証を切り分けてから、一致する本番用コンテナで再テストしてください。

性能を検証する場合は、プロファイリングなしで安定した時間帯の計測値を合否判定に使用してください。 因果関係を確認するには、条件を揃えた Nsight の A/B 比較を行います。

1. ディスパッチャー、ルーティング、CUDA グラフ、バッチ形状、並列化、ランタイムを固定します。
2. overlap_moe_expert_parallel_comm のみを切り替えます。最初の切り分けでは delay_wgrad_compute=false を維持してください。
3. カーネル時間の合計だけではなく、通信と計算の時間区間の和集合、および両者の共通部分を比較します。
4. 安定した状態のステップ時間、モデルの TFLOPS/GPU、損失が有限かどうか、スキップされた反復数／NaN の反復数、 ピーク割り当てメモリを報告します。

最終シグネチャ更新日：2026-08-03。