SuperAPI AI 工具集

nemo-mbridge-perf-expert-parallel-overlap

智能体与模型

查看项目

更新于 10-16 03:49

技能介绍

用途:在 Megatron-Bridge 中启用 MoE 专家并行(EP)通信重叠,让 token dispatch/combine 的 all-to-all 通信与专家 FFN 计算并行执行,从而隐藏通信开销。可选 delay_wgrad_compute 将专家权重梯度计算延后,与下一层前向重叠。

主要能力与用法:提供两种调度路径。alltoall 为默认、兼容性最广;flex 使用 DeepEP 或 HybridEP,在受支持的 GPU 上可获得更高重叠。基本启用方式为设置 cfg.comm_overlap.overlap_moe_expert_parallel_comm = True,保持 cfg.model.moe_shared_expert_overlap = False,并将 moe_token_dispatcher_type 设为 "alltoall" 或 "flex"。使用 flex 需调用 apply_flex_dispatcher_backend(cfg.model, moe_flex_dispatcher_backend="deepep")(或 "hybridep"),仅设置 backend 字段不会激活。仅在普通重叠路径验证可用后再启用 delay_wgrad_compute=True。

适用场景:EP > 1 的 MoE 模型、dispatch/combine 通信占步骤时间明显、有显存余量且以吞吐为调优目标时。资料给出 Qwen3 MoE 30B-A3B SFT 在 16×H100、EP=16、alltoall、BF16 下的隔离基准(迭代 3-8):无重叠 41.25s,EP 重叠 31.31s(1.317x),叠加 delay_wgrad 31.20s(1.322x);未显示 delayed wgrad 的独立显著收益。

限制与约束:需 expert_model_parallel_size > 1、num_moe_experts > 1、BF16 或 FP16、PyTorch >= 2.6.0;启用全量重算、moe_shared_expert_overlap 或仍在正确性搭建阶段时避免使用;PP > 1 需设置 virtual_pipeline_model_parallel_size;delay_wgrad_compute 以 EP 重叠为前提,并在 overlap_grad_reduce、gradient_accumulation_fusion、CUDA graph 等组合下对 TE 版本有额外要求;DeepEP 限 Ampere/Hopper/B200/B300,HybridEP 另支持 NVL72 的 GB200/GB300。EP 重叠与共享专家重叠互斥。验证可用 uv run python -m pytest tests/unit_tests/training/test_comm_overlap.py tests/unit_tests/training/test_deepep.py -q。公开配方常默认关闭 MoE 重叠,各模型族的端到端收益尚未在受控 Bridge 实验中全部测量。