nemo-mbridge-perf-moe-optimization-workflow

MoE 训练优化工作流

稳定文档：@docs/training/moe-optimization.md 卡片：@skills/nemo-mbridge-perf-moe-optimization-workflow/card.yaml 来源：使用 Megatron Core 对 MoE 模型进行可扩展训练

快速参考

从论文提出的三大瓶颈来思考：

• 内存瓶颈
• 通信瓶颈
• 计算和主机开销瓶颈

MoE 调优是一个迭代过程。解决一个瓶颈后，通常会暴露下一个瓶颈，因此最佳工作流是：先让模型适配，再扩展，然后进行性能分析，最后重新调优。

首次回答检查清单

对于 MoE 优化工作流相关提示词，请按以下顺序组织回答：

1. 适配：首先确保模型在内存允许范围内运行。使用能够容纳模型的最小模型并行度；优先使用选择性重计算，而不是完全重计算；只有在重计算和并行度都不足时才添加卸载；使用 --fake-init-process-group 对大型布局进行合理性检查。
2. 扩展：模型适配后，尽可能提高 DP；将高频通信限制在最快的互连范围内；多节点扩展时使用 PP 和 VPP；专家层优先使用 EP，而不是增加 TP；长上下文导致注意力内存占主导时，添加 CP。
3. 性能分析：识别主要瓶颈：内存、通信、主机开销或计算。
4. 重新调优：根据性能分析发现的瓶颈，调整分发器、重叠、FP8 模式、CUDA 图或重计算。
5. 写明 Parallel Folding 的准确网格：Attention: TP x CP x DP x PP 和 MoE: ETP x EP x EDP x PP。
6. 写明默认映射：安全启动时使用 alltoall；H100/B200 类系统使用 flex + deepep；GB200/GB300/NVL72 类系统使用 flex + hybridep；Hopper 使用 FP8 blockwise；Blackwell 使用 MXFP8；对于无丢弃 MoE，使用覆盖 attn、moe_router 和 moe_preprocess 的 TE 作用域 CUDA 图。

阶段 1：确保运行满足内存要求

先从能够可靠运行的配置开始，再追求吞吐量。

建议顺序：

1. 使用仍能容纳模型的最小模型并行度。
2. 启用选择性重计算；只有不够时才退回完全重计算。
3. 只有在重计算和并行度仍不足时才添加卸载。
4. 在单块 GPU 上使用 --fake-init-process-group 对大型并行布局进行合理性检查，避免浪费集群时间。

重计算指导

MoE 运行优先采用选择性重计算：

• 首选项：layernorm、core_attn、moe_act、mlp，或特定于模型的模块（shared_experts、mla_up_proj）
• 只有运行仍无法适配时才使用完全重计算
• 启用 CUDA 图后重新评估重计算，因为某些图作用域与完全重计算路径不太兼容

经验上，细粒度重计算通常能释放大部分所需内存，同时比全层重计算更接近不使用重计算时的吞吐量基线。

阶段 2：选择并行方式以扩展规模

优先顺序：

1. 模型适配后，尽可能提高 DP。
2. 尽可能将高频通信路径限制在高速互连范围内。
3. 多节点扩展时使用 PP，并在需要时使用 VPP。
4. 专家层优先使用 EP，而不是增加 TP。
5. 序列长度导致注意力内存占主导后，为长上下文添加 CP。

Parallel Folding

Parallel Folding 将注意力和 MoE 并行方式解耦，因此无需在两种布局之间选择单一折中方案：

Attention: TP × CP × DP × PP
MoE:       ETP × EP × EDP × PP

关键参数：

• --expert-model-parallel-size
• --expert-tensor-parallel-size

当注意力层需要一定的 TP 或 CP，但专家层受益于比稠密层所能承受的更高 EP 度数时，使用此功能。

第 3 阶段：分析主要瓶颈

瓶颈	表现	主要修复方法
内存	只有采用激进的全量重计算才能运行，或在预热期间发生 OOM	选择性重计算、FP8、卸载、更优的 PP 布局
通信	Nsight 显示大量 all-to-all 或 collective 区块	DeepEP 或 HybridEP、EP 重叠、DP/TP 重叠、更优的 PP 布局
主机开销	GPU 空闲间隙、受启动开销限制的跟踪、Python 开销	CUDA 图、--manual-gc、更高的 MBS、CPU 亲和性调优
计算	解决通信和主机问题后，SM 利用率仍然较低	分组 GEMM、融合优化、FP8、针对调度器的内核调优

调度器与重叠指南

将调度器选择作为解决瓶颈的手段，而不是首要调优项。

• moe_token_dispatcher_type="alltoall"：最稳妥的初始方案，适用于较小的 EP 规模
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="deepep"： 对 H100 和 B200 类型的部署来说，是很好的默认选择
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="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。最大的收益通常来自专家 GEMM 和其他高负载矩阵运算，而不是尝试量化 MoE 的每个小组件。

MoE 的 CUDA 图

对于无丢弃 MoE，从部分 TE 作用域图开始：

• attn
• moe_router
• moe_preprocess

这条路径通常能显著缩短步时，同时将动态专家计算留在图外。如果能看到启动开销，预期会有中等幅度的加速，但需预留数 GB 的额外内存，并确认形状保持静态。

仅对适合图的工作负载使用全迭代图，例如丢弃并填充，或严格控制的静态形状实验。

相关参考：

• @skills/nemo-mbridge-perf-cuda-graphs/SKILL.md
• @docs/training/cuda-graphs.md
• @docs/training/activation-recomputation.md

常见问题

1. 不要按错误的顺序优化：让模型适配硬件并选择合理的并行策略，比微优化更重要。

2. 平台会改变瓶颈：H100 级别的运行通常更受通信限制，而 GB200 或 GB300 则往往更早暴露 CPU 或启动开销。

3. FP8 MFU 可能显得异常低：切换精度模式时，除了 MFU，也要比较绝对吞吐量。

4. CUDA 图与重计算会相互影响：TE 作用域图通常与选择性重计算搭配，而不是与全量重计算搭配。

5. 大规模场景下，Parallel Folding 并非可选项：一旦注意力层和专家层明显需要不同布局，共享单一 TP 或 EP 方案就会成为双方的负担。