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 によってさらに制約が加わります。
• 実際には、オーバーラップを有効にする場合は、選択的再計算を組み合わせる方が安全です。

短時間計測時の注意

2026-05-18 に現行 main で実施した H100 x16 の Qwen3 30B-A3B モック事前学習のスモークテストでは、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

通信が顕在化しているノード間 alltoall の MoE 構成では、EP オーバーラップが有効な場合があることを示す結果として扱ってください。遅延 wgrad が単独で効果をもたらす証拠ではなく、融合パーミュテーション経路を検証するものでもありません。同じ構成で 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 実行で、常に効果があるとは限りません。

検証

初期化時にオーバーラップ関連のログメッセージを確認してください。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 \
  --config_variant v1 \
  --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 を追加して、オーバーラップの計測とランタイムスタックの検証を切り分けてから、対応する本番コンテナで再テストしてください。