SuperAPI AI 工具集

nemo-mbridge-perf-moe-comm-overlap

智能体与模型

查看项目

更新于 10-16 03:49

技能介绍

用途:在 Megatron Bridge 中开启 MoE 专家并行(EP)通信重叠,将 token dispatch/combine 通信与计算重叠,以降低暴露的通信开销、提升训练吞吐。适用于 EP>1、profile 中 dispatch 或 combine 时间明显、且训练已能正确跑通的调优阶段,不建议作为早期 bring-up 步骤。

主要能力:支持 dispatch/combine 重叠、flex dispatcher 后端(DeepEP 等)以及通过 delay_wgrad_compute 延迟专家 wgrad 计算以获得额外重叠。

启用配置: cfg.comm_overlap.overlap_moe_expert_parallel_comm = True cfg.comm_overlap.delay_wgrad_compute = True # 可选 cfg.model.moe_shared_expert_overlap = False # 使用 dispatch 重叠时必须关闭

前置条件: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)。

注意事项:仅设置 moe_flex_dispatcher_backend 不会激活 flex dispatch,必须同时设 moe_token_dispatcher_type = "flex"。全量重计算不适合该重叠路径,启用重叠时更安全的是选择性重计算;delay_wgrad_compute 在 CUDA graph 覆盖 attention 或 MoE router 时会有额外约束。共享专家重叠与 dispatch 重叠可能冲突,需关闭前者。多数公开 recipe 默认关闭 MoE 重叠,需显式 override 打开。收益依赖负载,通信已是 step 时间显著部分时最有效,并非所有小规模或轻载 EP 都获益。

验证与冒烟:初始化时查看重叠相关日志;comm_overlap.py 的校验会在前置条件不满足时报错。短跑性能可参考 scripts/performance/run_script.py 的命令形态,每次只改一个重叠开关;若融合 MoE permutation 失败,可加 model.moe_permute_fusion=false 以隔离重叠计时与运行时栈问题。资料中 2026-05-18 的 H100 x16 冒烟结果仅为方向性证据,不构成延迟 wgrad 单独收益或融合 permutation 路径的结论。

代码锚点:comm_overlap.py、flex_dispatcher_backend.py、config.py 及对应单元测试。