nemo-mbridge-perf-expert-parallel-overlap

MoE エキスパート並列オーバーラップスキル

参照

• 安定版ドキュメント：@docs/training/communication-overlap.md
• 構造化メタデータ：@skills/nemo-mbridge-perf-expert-parallel-overlap/card.yaml

概要

エキスパート並列（EP）オーバーラップは、エキスパート FFN の計算と並行してトークンのディスパッチ／結合の all-to-all 通信を実行し、そのコストを隠します。オプションで、エキスパートの重み勾配計算（delay_wgrad_compute）を遅延させることで、次の層のフォワード計算とオーバーラップさせ、さらにオーバーラップを増やせます。

Bridge は 2 種類のディスパッチャーパスに対応しています。

ディスパッチャー	バックエンド	使用する場面
alltoall	標準の MoE all-to-all	デフォルト。最も幅広く互換性があります
flex	DeepEP または HybridEP	Ampere／Hopper／Blackwell で、より大きなオーバーラップを得たい場合

すぐに判断するには

次の場合は EP オーバーラップを使用します。

• モデルが MoE で、EP > 1 の場合
• エキスパートのディスパッチ／結合の通信が、ステップ時間の中で無視できない割合を占める場合
• メモリに余裕があり、スループットを調整している場合

以下を優先します。

• 最初の導入では alltoall ディスパッチャー（より幅広い互換性）
• 対応 GPU で実行し、さらなる改善を目指す場合は flex + DeepEP／HybridEP

次の場合は EP オーバーラップを避けます。

• 全アクティベーションの再計算が有効になっている場合
• moe_shared_expert_overlap が有効になっている場合
• 正確性を確認しながら実行を立ち上げている段階の場合
• PyTorch < 2.6.0 の場合

期待される結果：

• all-to-all ディスパッチがプロファイル上で明確なボトルネックであれば、オーバーラップによって小幅から大幅な高速化が得られる可能性があります
• 実行規模が小さい場合、通信が少ない場合、または別の要因が実行時間の大半を占める場合は、改善がほとんど得られないことがあります

正確性を優先した alltoall ベンチマーク

通常の EP オーバーラップだけを切り分けてベンチマークする場合は、flex ディスパッチと遅延 wgrad を無効にします。測定した構成は、16 台の H100 GPU で実行した Qwen3 MoE 30B-A3B SFT です。 EP=16、alltoall、BF16、グローバルバッチサイズ 1024、CUDA グラフ無効、 moe_permute_fusion=false で、イテレーション 3～8 を測定しました。

通常のオーバーラップを使用する場合は、次のオーバーライドを使用します。

--cuda_graph_impl none \
--moe_flex_dispatcher_backend None \
--moe_a2a_overlap false \
comm_overlap.overlap_moe_expert_parallel_comm=true \
comm_overlap.delay_wgrad_compute=false \
model.moe_shared_expert_overlap=false

この切り分けテストでは --moe_a2a_overlap true を使用しないでください。パフォーマンスハーネスのヘルパーは overlap_moe_expert_parallel_comm と delay_wgrad_compute の両方を有効にするため、通常の EP オーバーラップだけを切り分けられません。

このベンチマークでの定常区間の測定時間：

ケース	定常時の平均	相対値
EP オーバーラップなし	41.25s	1.000x
EP オーバーラップあり	31.31s	1.317x
EP オーバーラップあり、delay_wgrad_compute も有効	31.20s	1.322x

これは、ノード間 all-to-all のこの構成で通常の EP オーバーラップを有効にする根拠となります。遅延 wgrad 単独で有意な改善が得られることを示すものではなく、実行時スタックの都合で無効にしていたため、MoE の融合パーミュテーションも検証していません。

有効化

alltoall ディスパッチャー

cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
cfg.comm_overlap.delay_wgrad_compute = False
cfg.model.moe_shared_expert_overlap = False

cfg.model.expert_model_parallel_size = 8
cfg.model.num_moe_experts = 64
cfg.model.moe_token_dispatcher_type = "alltoall"
cfg.model.bf16 = True
cfg.model.fp16 = False

通常のオーバーラップパスが動作することを確認し、追加の互換性要件を確認した後にのみ、delay_wgrad_compute=True を有効にしてください。

flex ディスパッチャー（DeepEP または HybridEP）

from megatron.bridge.training.flex_dispatcher_backend import apply_flex_dispatcher_backend

cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
cfg.comm_overlap.delay_wgrad_compute = True
cfg.model.moe_shared_expert_overlap = False

apply_flex_dispatcher_backend(cfg.model, moe_flex_dispatcher_backend="deepep")
# or: apply_flex_dispatcher_backend(cfg.model, moe_flex_dispatcher_backend="hybridep")

互換性と制約

• expert_model_parallel_size > 1
• num_moe_experts > 1
• moe_token_dispatcher_type は "alltoall" または "flex" である必要があります
• moe_shared_expert_overlap = False
• ベース精度は BF16 または FP16
• PyTorch >= 2.6.0
• PP > 1 の場合は、virtual_pipeline_model_parallel_size を設定する必要があります
• recompute_granularity != "full"、recompute_method = None、 recompute_num_layers = None
• mtp_num_layers は None または 1 である必要があります
• delay_wgrad_compute を使用するには、前提条件として overlap_moe_expert_parallel_comm が必要です
• overlap_grad_reduce と併用する delay_wgrad_compute には TE >= 2.7.0 が必要です
• gradient_accumulation_fusion と併用する delay_wgrad_compute には TE >= 2.7.0 が必要です
• CUDA グラフの attn スコープと delay_wgrad_compute の併用には、TE >= 2.12.0、 gradient_accumulation_fusion = True、およびアテンションバイアスなしが必要です
• DeepEP：Ampere、Hopper、B200、B300 GPU のみ
• HybridEP：Ampere、Hopper、B200、B300、NVL72 搭載の GB200／GB300

最小限の動作構成

cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
cfg.comm_overlap.delay_wgrad_compute = False
cfg.model.expert_model_parallel_size = 4
cfg.model.num_moe_experts = 64
cfg.model.moe_token_dispatcher_type = "alltoall"
cfg.model.moe_shared_expert_overlap = False
cfg.model.bf16 = True

これは、正しさを優先する出発点として使用してください。通常のオーバーラップ経路が動作することを確認してから、遅延 wgrad、flex ディスパッチ、CUDA グラフとの連携を追加してください。

最小限の実行可能なコマンド

Slurm の割り当て内で実行するパフォーマンスハーネスの例です。モデル、並列化、ディスパッチャー、ランタイムは固定し、2 つのオーバーラップ設定のみを変更してください。

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

通常の EP オーバーラップと遅延 wgrad を切り分ける際に、--moe_a2a_overlap true を使用しないでください。パフォーマンスハーネスのヘルパーは、overlap_moe_expert_parallel_comm と delay_wgrad_compute の両方を有効にします。

単体テストによる検証：

uv run python -m pytest \
  tests/unit_tests/training/test_comm_overlap.py -k "moe" \
  tests/unit_tests/training/test_deepep.py -q

検証

単体テスト

uv run python -m pytest \
  tests/unit_tests/training/test_comm_overlap.py \
  tests/unit_tests/training/test_deepep.py -q

ログの確認

EP オーバーラップを使用して正常に実行できた後、以下を確認してください。

1. CommOverlapConfig の確定時にアサーションエラーがないことを確認します
2. ログに出力される設定で、overlap_moe_expert_parallel_comm が True になっていることを確認します
3. flex ディスパッチャーを使用する場合は、ログで moe_token_dispatcher_type = "flex" と正しいバックエンドを確認します

成功基準

• 選択したディスパッチャーとオーバーラップ設定で構成の検証に合格する
• 学習がハングやアサーション失敗なく完了する
• 対象ワークロードのスループットが向上するか、少なくとも低下しない
• 損失の推移がベースラインと一致する（オーバーラップは収束に影響しないはずです）

コードの参照箇所

Bridge のオーバーラップ検証

if self.user_comm_overlap_cfg.overlap_moe_expert_parallel_comm is True:
    assert model_cfg.expert_model_parallel_size > 1, ...
    assert model_cfg.num_moe_experts > 1, ...
    assert model_cfg.moe_token_dispatcher_type in ["alltoall", "flex"], ...
    assert model_cfg.bf16 or model_cfg.fp16, ...
    assert is_torch_min_version("2.6.0"), ...
    # ... PP + VPP check, recompute checks, shared_expert_overlap check ...

遅延 wgrad の検証

if self.user_comm_overlap_cfg.delay_wgrad_compute is True:
    # TE version checks for overlap_grad_reduce and gradient_accumulation_fusion
    # CUDA graph scope validations for delayed wgrad
    assert overlap_moe_expert_parallel_comm, ...

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

def apply_flex_dispatcher_backend(...):
    # GPU architecture check for DeepEP / HybridEP
    model_config.moe_token_dispatcher_type = "flex"
    model_config.moe_flex_dispatcher_backend = moe_flex_dispatcher_backend
    model_config.moe_shared_expert_overlap = False

パフォーマンスハーネスの設定上書き

def _set_moe_a2a_overlap_overrides(recipe, moe_a2a_overlap=False):
    if moe_a2a_overlap:
        recipe.comm_overlap.overlap_moe_expert_parallel_comm = True
        recipe.comm_overlap.delay_wgrad_compute = True
        recipe.model.moe_shared_expert_overlap = False

テスト

ファイル	対象範囲
tests/unit_tests/training/test_comm_overlap.py	EP オーバーラップの検証、遅延 wgrad、CUDA グラフと wgrad の相互作用
tests/unit_tests/training/test_deepep.py	DeepEP/HybridEP ヘルパーの有効化と GPU の適用条件

問題の診断

症状	考えられる原因	確認方法	修正方法
assert expert_model_parallel_size > 1	EP が設定されていない	expert_model_parallel_size を確認します	EP を 1 より大きく設定します
assert moe_token_dispatcher_type	ディスパッチャーが正しくない	ディスパッチャーの種類を確認します	"alltoall" または "flex" を使用します
BF16/FP16 に関する assert	精度が正しくない	bf16 と fp16 を確認します	bf16 = True に設定します
学習中にハングする	PyTorch が 2.6 未満	PyTorch のバージョンを確認します	>= 2.6.0 にアップグレードします
assert virtual_pipeline_model_parallel_size	VPP なしで PP > 1	PP と VPP の設定を確認します	PP > 1 の場合は VPP を設定します
assert recompute_granularity	完全再計算が有効	再計算の設定を確認します	完全再計算を無効にします
assert overlap_moe_expert_parallel_comm required	EP オーバーラップなしで遅延 wgrad を使用している	オーバーラップなしで delay_wgrad_compute が設定されているか確認します	先に EP オーバーラップを有効にします
assert gradient_accumulation_fusion	CUDA グラフと遅延 wgrad の組み合わせ	グラフの適用範囲と wgrad の設定を確認します	gradient_accumulation_fusion を有効にします
attention bias に関する assert	CUDA グラフの attention + 遅延 wgrad + bias	add_bias_linear / add_qkv_bias を確認します	attention bias を無効にします
flex ディスパッチャーでスループットが向上しない	apply_flex_dispatcher_backend が呼び出されていない	ログの moe_token_dispatcher_type を確認します	apply_flex_dispatcher_backend(...) を呼び出します
DeepEP/HybridEP が警告なくスキップされる	GPU が未対応	警告ログを確認します	Ampere/Hopper/Blackwell で実行します

既知の制限事項

• moe_flex_dispatcher_backend を設定するだけでは flex ディスパッチは有効になりません。apply_flex_dispatcher_backend(...) を呼び出す必要があります。
• 公開されているレシピは保守的なことが多く、MoE オーバーラップはデフォルトで無効になっています。
• すべてのモデルファミリーについて、Bridge の管理された実験でエンドツーエンドのスループット向上が測定されたわけではありません。コードの検証は、単一の普遍的な性能主張よりも確かなものです。
• MoE オーバーラップと共有エキスパートのオーバーラップは同時に使用できません。
• CUDA グラフと遅延 wgrad の組み合わせは複数の制約がある経路であり、TE のバージョンと適用範囲を慎重に検証する必要があります。