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

安定版ドキュメント: @docs/training/communication-overlap.md カード: @skills/perf-expert-parallel-overlap/card.yaml

参考資料

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

概要

エキスパート並列（EP）オーバーラップは、トークンディスパッチ/結合 all-to-all 通信のコストを、エキスパート FFN コンピュートと同時に実行することで隠します。オプションで、遅延エキスパート重み勾配計算（delay_wgrad_compute）は wgrad を次のレイヤーの順伝播と重なるようにすることで、追加のオーバーラップを提供します。

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 ディスパッチャ

cfg.comm_overlap.overlap_moe_expert_parallel_comm = True
cfg.comm_overlap.delay_wgrad_compute = True
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

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")
# または: 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 が必要です
• delay_wgrad_compute と overlap_grad_reduce の組み合わせには TE >= 2.7.0 が必要です
• delay_wgrad_compute と gradient_accumulation_fusion の組み合わせには 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、GB200/GB300 with NVL72

最小限の作動設定

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 グラフの相互作用は、プレーンなオーバーラップパスが動作することが確認できた後でのみ追加します。

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

パフォーマンスハーネスの例：

python scripts/performance/setup_experiment.py \
  --model qwen3-30b-a3b \
  --moe_a2a_overlap \
  --num_nodes 2 \
  --gpus_per_node 8 \
  --max_steps 20

単体テスト検証：

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 チェック、再計算チェック、shared_expert_overlap チェック ...

遅延 wgrad 検証

if self.user_comm_overlap_cfg.delay_wgrad_compute is True:
    # overlap_grad_reduce と gradient_accumulation_fusion の TE バージョンチェック
    # 遅延 wgrad の CUDA グラフスコープ検証
    assert overlap_moe_expert_parallel_comm, ...

Flex ディスパッチャアクティベーション

def apply_flex_dispatcher_backend(...):
    # DeepEP / HybridEP の GPU アーキテクチャチェック
    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" を使用
assert on BF16/FP16	不正な精度	bf16 と fp16 を確認	bf16 = True を設定
トレーニング中にハング	PyTorch < 2.6	PyTorch バージョンを確認	>= 2.6.0 にアップグレード
assert virtual_pipeline_model_parallel_size	PP > 1 で VPP なし	PP と VPP 設定を確認	PP > 1 の場合は VPP を設定
assert recompute_granularity	完全な再計算が有効化	再計算設定を確認	完全な再計算を無効化
assert overlap_moe_expert_parallel_comm required	オーバーラップなしの遅延 wgrad	delay_wgrad_compute をオーバーラップなしで確認	最初に EP オーバーラップを有効化
assert gradient_accumulation_fusion	CUDA グラフ + 遅延 wgrad	グラフスコープ + wgrad 設定を確認	gradient_accumulation_fusion を有効化
assert on アテンションバイアス	CUDA グラフ attn + 遅延 wgrad + バイアス	add_bias_linear / add_qkv_bias を確認	アテンションバイアスを無効化
flex ディスパッチャから性能向上がない	apply_flex_dispatcher_backend が呼ばれていない	ログで moe_token_dispatcher_type を確認	apply_flex_dispatcher_backend(...) を呼び出し
DeepEP/HybridEP がサイレントにスキップ	サポートされていない GPU	警告ログを確認	Ampere/Hopper/Blackwell で実行

既知の制限事項

• moe_flex_dispatcher_backend を単独で設定してもフレックスディスパッチはアクティベートされません。apply_flex_dispatcher_backend(...) を呼び出す必要があります。
• パブリックレシピはしばしば保守的で、MoE オーバーラップをデフォルトで無効化したままにしています。
• エンドツーエンドのスループット向上は、Bridge の制御された実験ですべてのモデルファミリについてまだ測定されていません。コード検証は単一の普遍的な性能主張より強力です。
• MoE オーバーラップと共有エキスパートオーバーラップは相互排他的です。
• CUDA グラフと遅延 wgrad は、TE バージョンとスコープの慎重な検証が必要な多制約パスです。