用途:为 Megatron Bridge 中的 MoE 训练提供系统化优化工作流,依据 Megatron-Core MoE 论文的“三墙”框架(内存墙、通信墙、计算与主机开销墙)组织调优思路。核心原则是先适配、再扩展、后剖析、最后再调优,因为解决一堵墙通常会暴露下一堵墙。
主要能力与流程:
- 适配:先让模型内存可行。使用能容纳的最小模型并行度,优先选择性重计算而非全量重计算,仅在重计算与并行度仍不足时加入 offloading,并用 --fake-init-process-group 在单 GPU 上校验大规模并行布局。选择性重计算首选 layernorm、core_attn、moe_act、mlp 或模型特定模块(shared_experts、mla_up_proj)。
- 扩展:模型适配后最大化 DP,把热通信留在最快互连内,多节点用 PP 加 VPP,专家层优先 EP 而非额外 TP,长上下文使注意力内存占主导时加入 CP。并行折叠解耦注意力与 MoE 并行:Attention 为 TP × CP × DP × PP,MoE 为 ETP × EP × EDP × PP,关键参数为 --expert-model-parallel-size 与 --expert-tensor-parallel-size。
- 剖析:识别主导瓶颈是内存、通信、主机开销还是计算,并给出对应修复方向。
- 再调优:依据剖析结果调整调度器、重叠、FP8 模式、CUDA 图或重计算。
调度器与重叠:alltoall 为最安全的启动路径;flex + deepep 适合 H100/B200 类系统;flex + hybridep 适合 GB200/GB300/NVL72 系统。若剖析中出现 all-to-all 路径,可结合 --overlap-moe-expert-parallel-comm、--overlap-grad-reduce、--tp-comm-overlap。
FP8 配方:Hopper 用 FP8 blockwise,Blackwell 用 MXFP8,Blackwell 速度优先探索可在 BF16 或 FP8 稳定后尝试 NVFP4;路由器保持 FP32。
CUDA 图:dropless MoE 从 TE 作用域的部分图开始,覆盖 attn、moe_router、moe_preprocess,可带来步时收益,但需额外数 GB 内存并确保形状静态;全迭代图仅用于 drop-and-pad 等图友好负载。
限制与注意:不要颠倒优化顺序;平台改变限制墙(H100 类常更偏通信受限,GB200/GB300 更早暴露 CPU 或启动开销);FP8 的 MFU 可能偏低,需同时比较绝对吞吐;CUDA 图与重计算存在交互,TE 作用域图通常搭配选择性重计算;大规模下并行折叠并非可选。具体命令、参数与平台细节以官方文档为准,资料未提供的部分暂未提供。