SuperAPI AI 工具集

nemo-mbridge-perf-moe-comm-overlap

编程与开发

查看项目

更新于 10-16 02:33

技能介绍

用途:在 Megatron Bridge 中启用 MoE 专家并行通信重叠,让 token dispatch/combine 与计算重叠,从而提升训练吞吐。适用于 EP>1、profile 中 dispatch 或 combine 时间明显、且训练已正确、进入吞吐调优阶段的场景;不建议作为早期 bring-up 步骤。

主要能力:支持 dispatch/combine 重叠、flex dispatcher 后端、专家 wgrad 延迟调度(delay_wgrad_compute)。

启用方式: 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)。

Flex dispatcher:仅设置 moe_flex_dispatcher_backend 不会激活 flex dispatch,必须同时设置 moe_token_dispatcher_type = "flex"。

重计算与 CUDA Graph:完整重计算不适合与重叠路径搭配;delay_wgrad_compute 在 CUDA-graph 作用域包含 attention 或 MoE-router 时还有额外约束。实践中启用重叠时选择性重计算更安全。

验证:初始化时查看重叠相关日志;comm_overlap.py 的校验会在前提不满足时报错,干净启动即表示已激活。短程性能冒烟可参考 run_script.py 命令,一次只调整一个重叠开关。

限制:多数公开 recipe 默认关闭 MoE 重叠,需显式覆盖开启;收益依赖工作负载,通信占比高时更明显,小规模或低负载 EP 不保证提升。资料中的 H100 x16 短程结果为方向性证据,非发布级结论,也未验证 fused permutation 路径。若 bring-up 时 fused MoE permutation 失败,可加 model.moe_permute_fusion=false 隔离重叠计时与运行时栈验证。