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 调优是迭代过程，因此请按以下顺序进行：

freeze the measurement contract -> fit -> scale -> profile -> retune -> validate

首要回答检查清单

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

1. 固定测量约定：记录确切的模型和任务、硬件及拓扑、容器和提交版本、数据及路由语义、精度、序列和批次形状、并行策略、图捕获范围，以及稳态指标窗口。将每个候选方案标记为与训练等效或仅用于基准测试。
2. 适配：首先确保模型在内存上可行。采用能够容纳模型的最小模型并行度；优先选择性重计算，而非完整重计算；只有在重计算和并行策略仍不够时才添加卸载；使用 --fake-init-process-group 检查大型布局是否合理。
3. 扩展：模型适配后，尽可能提高 DP；将高频通信限制在最快的互连范围内；多节点扩展时使用 PP 加 VPP；专家层优先使用 EP，而不是增加 TP；当长上下文导致注意力内存占主导时，添加 CP。
4. 分析性能：确定主要瓶颈：内存、通信、主机开销或计算。
5. 重新调优：根据性能分析确定的瓶颈，每次只更改一个变量。Dispatcher、重叠、低精度、CUDA 图和重计算都是候选项，而非硬件默认设置。
6. 验证：先用匹配的短时测试淘汰候选方案，然后让胜出方案至少运行 50 步。确认实际运行了所请求的后端或图重放；对明确声明的预热后窗口计时，并报告损失健康状况、跳过的迭代／NaN 迭代、内存、步时和模型 TFLOPS/GPU。
7. 明确列出 Parallel Folding 网格：Attention: TP x CP x DP x PP 和 MoE: ETP x EP x EDP x PP。
8. 使用 alltoall 进行安全的初始验证；如果软件包和目标拓扑支持，再对 flex + deepep 和 flex + hybridep 进行 A/B 测试。从 BF16 和急切执行开始；只有在性能分析证明合理后，才引入低精度或范围最小且有用的 CUDA 图捕获。

阶段 0：固定测量约定

除非正在测试的唯一变量就是以下某项，否则比较时必须保持这些条件不变：

• Bridge、MCore、Transformer Engine、容器、CUDA 和 NCCL 版本
• GPU 数量、SKU、节点拓扑，以及启动器／环境设置
• 模型、任务、数据路径、序列长度、MBS、GBS 和优化器设置
• TP、PP、VPP、CP、EP、ETP 和 DP 布局
• 路由语义、精度、重计算、dispatcher、重叠和图捕获范围
• 预热和稳态计时窗口

区分两种验收类别：

• 与训练等效的更改会保留预期的路由、损失、数据、优化器以及检查点／恢复行为。
• 仅用于基准测试的更改（例如强制负载均衡）适用于受控的内核研究，但不能证明生产训练的正确性或收敛性。

阶段 1：确保运行在内存上可行

先使用能够稳定运行的配置，再追求吞吐量。

建议顺序：

1. 使用仍能容纳模型的最小模型并行度。
2. 先启用选择性重计算；若仍不够，再回退到完整重计算。
3. 只有在重计算和并行策略仍不够时才添加卸载。
4. 在消耗集群资源前，使用 --fake-init-process-group 在单张 GPU 上检查大型并行布局是否合理。

重新计算指导

对于 MoE 运行，优先采用选择性重新计算：

• 首选项：layernorm、core_attn、moe_act、mlp，或 模型专属模块（shared_experts、mla_up_proj）
• 仅当运行仍无法容纳时才使用完整重新计算
• 启用 CUDA graphs 后，重新评估重新计算设置，因为某些图作用域与 完整重新计算路径并不兼容

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

阶段 2：选择适合扩展的并行方式

优先级顺序：

1. 模型适配后，尽可能增大 DP。
2. 尽可能将高频通信路径保留在高速互连内。
3. 多节点扩展时使用 PP，并在需要时使用 VPP。
4. 对专家层，优先选择 EP，而非增加 TP。
5. 当序列长度使注意力内存占主导时，添加 CP 以支持长上下文。

并行折叠

并行折叠将注意力和 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 或集体通信区段	DeepEP 或 HybridEP、EP 重叠、DP/TP 重叠、优化 PP 布局
主机开销	GPU 空闲间隙、受启动开销限制的跟踪、Python 开销	CUDA graphs、--manual-gc、增大 MBS、调整 CPU 亲和性
计算	解决通信和主机问题后，SM 利用率仍然偏低	分组 GEMM、融合优化、FP8、针对分发器的内核调优

分析重叠，避免误读

使用未进行性能分析的稳定迭代作为验收指标，并使用匹配的性能分析来解释原因：

1. 每次只更改一个重叠或分发器变量；保持路由、图作用域、 并行方式、批次形状和运行时不变。
2. 为通信内核和计算内核构建区间并集，再测量其交集，以量化隐藏的通信。
3. 不要累加内核时长并将结果视为墙钟时间。即使可见的 GPU 活动区间并集和端到端步进时间下降， 并发内核也可能因 SM 或带宽争用而运行更久。
4. 使用分发/合并 NVTX 范围、稳定步进时间、模型 TFLOPS/GPU、损失有限性、跳过/NaN 计数和峰值内存来交叉验证跟踪结果。

在受控的 16×H100 Qwen3 30B-A3B HybridEP 运行中，普通 EP 重叠将 GEMM/注意力隐藏的通信比例 从 0.11% 提高到 36.55%。未进行性能分析的步进时间从 24.7138s 降至 20.9920s，吞吐量则从 244.039 升至 287.305 model TFLOPS/GPU。delay_wgrad_compute 仍处于禁用状态。

阶段 4：根据证据重新调优

选择针对已分析瓶颈的最小候选方案，每次只更改一个变量。

调度器与重叠计算指南

选择调度器是为了解决瓶颈，而不是查硬件对照表。

• moe_token_dispatcher_type="alltoall"：最稳妥的初始方案，适用于较小的 EP 规模
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="deepep"： 在已安装 DeepEP 且通信路径可见时可考虑
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="hybridep"： 在已安装 HybridEP 的 NVL8 和 NVL72 系统上均可考虑，但需视拓扑而定

对于典型的 16×H100 Qwen3 30B-A3B 配置，HybridEP 加普通 EP 重叠计算是当前实测表现最佳的方案；而典型的 256×H100 Qwen3 235B 配置则使用标准 alltoall 加重叠计算。请在目标容器中基准测试后端兼容性和吞吐量；GPU 名称和 EP 度数都不能单独决定最佳方案。

如果性能分析结果显示全对全路径是瓶颈，请结合以下选项调整调度器：

• --overlap-moe-expert-parallel-comm
• --overlap-grad-reduce
• --tp-comm-overlap

首先分别测试普通 EP 重叠、共享专家重叠和延迟权重梯度计算。即使某个组件在另一种模型上有效，组合使用也可能导致性能回退。

低精度候选方案矩阵

先建立经过验证的 BF16 基线。硬件能力只决定哪些低精度候选方案可用，并不保证能提升速度。

平台	BF16 稳定后的候选方案
Hopper	目标技术栈支持的逐张量、当前缩放或分块 FP8
Blackwell	MXFP8 或其他受支持的 FP8 方案
Blackwell，优先追求速度	BF16/FP8 路径稳定后尝试 NVFP4

保持路由器使用 FP32。最大的性能提升通常来自专家 GEMM 和其他密集型矩阵计算，而不是尝试量化 MoE 的每个小组件。必须通过日志或跟踪确认预期的内核已运行，并根据端到端稳定迭代耗时评估候选方案，而非理论峰值 FLOPS。

MoE 的 CUDA Graph

只有在性能分析显示主机或内核启动间隙显著时，才使用 CUDA Graph。对于无丢弃的 MoE，请从范围最小的部分 TE 范围图候选项开始：

• moe_router
• moe_preprocess

只有在模型支持 attn，且它在相同匹配的技术栈中能提升性能时，才将其加入。成功捕获并不代表性能有所提升；调度器、重叠计算或精度发生变化后，图带来的性能优势也可能消失。

此路径会将动态专家计算保留在图外。预留额外内存，确认形状保持静态，并确认发生了重放而不只是捕获；只计时捕获后的迭代。

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

相关参考资料：

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

阶段 5：验证并整理证据

在工作负载允许的情况下，使用预热后的 6–12 次迭代进行低成本初筛。对选定的候选方案，至少运行 50 步，并报告固定的稳定窗口，例如第 41–50 步。最终证据包应包含：

• 精确的命令／配置差异、提交记录、容器、硬件和拓扑
• 声明的路由／数据语义，以及“训练等效”或“仅用于基准测试”的标签
• 预期的调度器、精度内核、重叠计算和图重放均已启用的证明
• 同一未进行性能分析的稳定窗口内的步耗时和模型 TFLOPS/GPU
• 有限损失值、跳过／NaN 计数、峰值内存，以及生产路径要求时的检查点／优化器状态验证
• 能证明所声称因果机制的匹配 A/B 性能分析证据

不要将最终多项改动方案的总收益归因于更早进行的单项 A/B 测试。例如，Qwen3 重叠实验单独测得 TFLOPS/GPU 从 244.039 提升至 287.305；之后的典型配置在进一步调整 HybridEP 后达到了 299.352。它们回答的是不同的问题。

易踩问题

1. 不要以错误的顺序优化：拟合模型并选择合理的并行度，比微优化更重要。

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

3. FP8 MFU 可能低得具有误导性：切换精度模式时，除了 MFU，也要比较绝对吞吐量。

4. CUDA 图与重计算会相互影响：限定于 TE 的图通常与选择性重计算搭配使用，而不是与全面重计算搭配。

5. 大规模场景下，Parallel Folding 并非可选项：一旦注意力层和专家层需要明显不同的布局，采用单一共享 TP 或 EP 方案就会让两者都付出代价。

6. 内核时间总和不等于暴露时间：验证重叠情况时，应使用区间并集以及通信与计算的交集。

7. 仅用于基准测试的语义不代表生产环境验收：强制路由、合成数据或禁用优化器／检查点路径等情况必须明确披露，并与训练等效结果分开验证。

8. 启用功能需要证据：如果后端可能回退、图可能捕获成功却没有帮助，或低精度方案可能未能触发预期内核，那么仅有配置转储是不够的。

最后一次签名刷新：2026-08-03。