SuperAPI AI 工具集

nemo-mbridge-perf-moe-optimization-workflow

办公与效率

查看项目

更新于 10-16 03:49

技能介绍

该技能提供一套系统化的 MoE 训练优化流程,依据 Megatron-Core MoE 论文的“三墙”框架:内存墙、通信墙、计算与主机开销墙。核心思路是先能跑起来,再扩规模,然后做性能剖析,最后重新调参。

使用顺序:第一步让运行内存可行。选用能装下的最小模型并行度,优先选择性重计算(如 layernorm、core_attn、moe_act、mlp 等),不足时再加 offloading,可用 --fake-init-process-group 先验证大规模并行布局。第二步选并行以扩展规模:模型装下后最大化 DP,让热通信留在最快互连内,多机用 PP 加 VPP,专家层优先 EP 而非额外 TP,长上下文注意力内存占主导时加 CP。并行折叠的网格为:注意力 TP × CP × DP × PP,MoE 为 ETP × EP × EDP × PP,关键参数为 --expert-model-parallel-size 与 --expert-tensor-parallel-size。第三步剖析主要瓶颈:内存、通信、主机开销或计算,再针对性调整调度器、重叠、FP8 模式、CUDA graphs 或重计算。

调度器建议:alltoall 适合安全启动和较小 EP;flex 加 deepep 适合 H100/B200;flex 加 hybridep 适合 GB200/GB300 NVL72。可结合 --overlap-moe-expert-parallel-comm、--overlap-grad-reduce、--tp-comm-overlap。FP8 方面,Hopper 用 FP8 blockwise,Blackwell 用 MXFP8,追求速度可在 BF16/FP8 稳定后尝试 NVFP4,路由保持 FP32。无丢弃 MoE 的 CUDA graphs 建议从 attn、moe_router、moe_preprocess 的 TE 范围图开始,需额外数 GB 内存并确保静态形状。

限制与注意事项:不要颠倒优化顺序;平台不同瓶颈墙不同;FP8 下 MFU 可能误导,应同时比较绝对吞吐;CUDA graphs 与重计算会相互影响,TE 范围图通常搭配选择性重计算而非全量重计算;大规模下并行折叠几乎是必需的。具体性能数字与硬件效果暂未提供。