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 パーミュテーションの妥当性も検証していません。

HybridEP の本番形状ベンチマーク

2026-07-25 に実施した、Qwen3 30B-A3B の事前学習における対照比較では、本番用 HybridEP パスで通常の EP オーバーラップが検証されました。

Hardware: 16×H100
Precision: BF16
Sequence: 4096
Parallelism: TP1 / PP1 / CP1 / EP16
Batch: MBS1 / GBS1024
Routing: force balance
Dispatcher: flex + HybridEP
CUDA graph: Transformer Engine scopes moe_router + moe_preprocess
Delayed wgrad: disabled

ケース	安定区間	ステップ時間	モデル TFLOPS／GPU
オーバーラップ無効	イテレーション 5～20	24.7138s	244.039
オーバーラップ有効、探索実行	イテレーション 5～20	21.0725s	286.208
オーバーラップ有効、独立検証	イテレーション 41～50	20.9920s	287.305

独立した実行では、再現したベースラインと比べてステップ時間が 15.059% 短縮し、スループットが 17.729% 向上しました。損失は有限値で、スキップされたイテレーションと NaN のイテレーションはいずれもゼロでした。また、rank-0 のピーク割り当てメモリは 62.166 GiB でした。

条件を一致させた Nsight Systems の比較では、各ケースで rank-0 のカーネル 463,348 件が記録されました。オーバーラップを有効にすると、GEMM およびアテンションと並行する通信が 9.079ms（通信時間の 0.11%）から 3,958.997ms（36.55%）に増加しました。GPU アクティブ区間の和集合は 22.821s から 21.221s に減少しました。

この結果は、オーバーラップの仕組みを裏付けるものとして使用してください。普遍的な高速化を保証するものではありません。ディスパッチャー、グラフのスコープ、ルーティング、並列化、バッチ形状、実行環境を固定し、通常の EP オーバーラップのみを変更しました。

有効化

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 = False
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")

まずプレーンな EP オーバーラップをベンチマークしてください。delay_wgrad_compute=True は、その CUDA グラフおよび TE との互換性制約が満たされた後に、独立したフォローアップの A/B としてのみ有効にしてください。

互換性と制約

• 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 \
  --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" と正しいバックエンドを確認します

成功基準

• 選択したディスパッチャーとオーバーラップ設定に対して設定検証が通ること
• トレーニング実行がハングやアサーション失敗なしで完了すること
• 対象ワークロードでスループットが向上する、または少なくとも劣化しないこと
• 損失の軌跡がベースラインと一致すること (オーバーラップは収束に影響を与えないはずです)

プロファイルの解釈

スループットの受け入れ結果には、プロファイリングされていない定常ウィンドウを使用してください。メカニズムを説明するには、対応するプロファイルを使用してください:

1. ディスパッチャー、ルーティング、グラフスコープ、バッチ形状、並列レイアウト、ランタイムを固定してください。
2. プレーンな EP オーバーラップのみを切り替えながら、同じランクと同じ定常イテレーションをキャプチャしてください。
3. 通信カーネルと計算カーネルの区間の和集合を構築し、それらの積集合を測定してください。
4. カーネル継続時間の合計をウォールタイムとして使用しないでください。同時実行カーネルは、公開時間が減少しても、SM または帯域幅の競合により長く実行されることがあります。
5. 区間の結果を、dispatch/combine の NVTX レンジ、最終ステップ時間、損失の有限性、スキップ/NaN の数、およびピークメモリで裏付けてください。

コードアンカー

ブリッジオーバーラップ検証

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 ゲーティング

障害の診断

症状	考えられる原因	確認方法	対処方法
expert_model_parallel_size > 1 のアサーション	EP が設定されていません	expert_model_parallel_size を確認します	EP を 1 より大きく設定します
moe_token_dispatcher_type のアサーション	ディスパッチャーが正しくありません	ディスパッチャーの種類を確認します	"alltoall" または "flex" を使用します
BF16/FP16 のアサーション	精度設定が正しくありません	bf16 と fp16 を確認します	bf16 = True に設定します
トレーニング中にハングする	PyTorch が 2.6 未満です	PyTorch のバージョンを確認します	>= 2.6.0 にアップグレードします
virtual_pipeline_model_parallel_size のアサーション	VPP なしで PP が 1 より大きく設定されています	PP と VPP の設定を確認します	PP が 1 より大きい場合は VPP を設定します
recompute_granularity のアサーション	全再計算が有効になっています	再計算の設定を確認します	全再計算を無効にします
overlap_moe_expert_parallel_comm required のアサーション	EP のオーバーラップなしで wgrad が遅延されています	オーバーラップなしで delay_wgrad_compute が設定されていないか確認します	まず EP のオーバーラップを有効にします
gradient_accumulation_fusion のアサーション	CUDA グラフと遅延 wgrad が併用されています	グラフのスコープと wgrad の設定を確認します	gradient_accumulation_fusion を有効にします
アテンションバイアスのアサーション	CUDA グラフのアテンション、遅延 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 を設定するだけでは flex dispatch は有効になりません。 apply_flex_dispatcher_backend(...) を呼び出す必要があります。
• 公開されているレシピは保守的な設定であることが多く、デフォルトでは MoE オーバーラップが無効になっています。
• Qwen3 30B-A3B の HybridEP H100 の構成 1 つについて、条件を制御したエンドツーエンドの検証結果とプロファイルの証拠があります。別のモデル、ディスパッチャー、トポロジー、精度、またはバッチ構成に一般化する前に、条件を一致させた A/B テストを再実施してください。
• MoE オーバーラップと共有エキスパートのオーバーラップは相互に排他的です。
• CUDA グラフと遅延 wgrad を組み合わせる場合は、複数の制約があるため、TE のバージョンとスコープを慎重に検証する必要があります。

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