⚡

MoE性能优化工作流Skill nemo-mbridge-perf-moe-optimization-workflow

本技能提供一套基于证据的MoE（混合专家）模型训练性能优化工作流，适用于Megatron Bridge环境。涵盖测量契约、三堵墙（内存/通信/计算）诊断、并行折叠、性能剖析、Dispatcher与CUDA图调优及验证，用于解决MoE吞吐量调优与性能回归诊断。关键词：MoE、Megatron、性能优化、内存墙、通信墙、计算墙、并行折叠、DeepEP、HybridEP、CUDA图、FP8、A/B测试。

分布式训练调优 0 次安装 0 次浏览 更新于 9/7/2026
名称 nemo-mbridge-perf-moe-optimization-workflow
描述 基于证据的MoE性能优化工作流，适用于Megatron Bridge。涵盖测量契约、三堵墙框架、并行折叠、性能剖析、匹配A/B调优以及最终验证。
开源协议 Apache-2.0 when_to_use: 完整的MoE吞吐量调优，或在提交或配置更改后诊断MoE吞吐量回归；‘优化MoE吞吐量’、‘MoE性能调优’、‘三堵墙’、‘内存墙’、‘通信墙’、‘计算墙’。

MoE训练优化工作流

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

快速参考

先从论文的“三堵墙”开始：

  • 内存墙
  • 通信墙
  • 计算效率墙

对于操作诊断，将计算效率墙拆分为计算和主机/启动瓶颈。它们需要不同的证据和不同的修复。MoE调优是迭代的，因此使用以下顺序：

冻结测量契约 -> 适配 -> 扩展 -> 剖析 -> 重新调优 -> 验证

首要答案检查清单

对于MoE优化工作流提示，请按以下顺序呈现响应：

  1. 冻结测量契约：记录确切的模型和任务、硬件和拓扑、容器和提交、数据和路由语义、精度、序列和批次形状、并行性、图范围，以及稳态指标窗口。将每个候选项标记为训练等价或仅基准测试。
  2. 适配：首先使模型在内存上可行。使用能容纳的最小模型并行度，优先选择性重新计算而不是完全重新计算，仅当重新计算和并行不足以缓解时添加卸载，并使用 --fake-init-process-group 对大型布局进行健全性检查。
  3. 扩展：模型适配后最大化DP，将热点通信保持在最快互连内，使用PP加VPP进行多节点扩展，优先使用EP而不是额外TP来处理专家层，并在长上下文使注意力内存主导时添加CP。
  4. 剖析：识别主导瓶颈：内存、通信、主机开销或计算。
  5. 重新调优：基于剖析出的瓶颈一次只改变一个变量。调度器、重叠、较低精度、CUDA图、重新计算都是候选，而非硬件默认值。
  6. 验证：使用短期匹配筛选来拒绝候选，然后让胜者至少运行50步。验证请求的后端或图重放确实运行，计时声明的预热后窗口，并报告损失健康状况、跳过/NaN迭代、内存、步进时间和模型TFLOPS/GPU。
  7. 包含精确的并行折叠网格：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布局
  • 路由语义、精度、重新计算、调度器、重叠和图范围
  • 预热和稳态计时窗口

区分两种接受类别：

  • 训练等价更改保留预期的路由、损失、数据、优化器和检查点/恢复行为。
  • 仅基准更改（如强制负载均衡）对于受控内核研究很有用，但无法建立生产训练的准确性或收敛性。

阶段1：使运行内存可行

在追求吞吐量之前，从一个能可靠容纳的配置开始。

推荐顺序：

  1. 使用仍然容纳的最小模型并行度。
  2. 在回退到完全重新计算之前启用选择性重新计算。
  3. 仅当重新计算和并行度仍然不足时添加卸载。
  4. 使用 --fake-init-process-group 在单个GPU上健全性检查大型并行布局，然后才消耗集群时间。

重新计算指南

对于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用于长上下文。

并行折叠

并行折叠解耦注意力和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图、--manual-gc、更高MBS、CPU亲和性调优
计算 在解决通信和主机问题后SM利用率低 分组GEMM、融合工作、FP8、调度器特定内核调优

剖析重叠而不误读

使用未剖析的稳态迭代作为验收指标，并使用匹配的剖析进行因果解释：

  1. 一次更改一个重叠或调度器变量；保持路由、图范围、并行度、批次形状和运行时固定。
  2. 为通信内核和计算内核构建区间并集，然后测量交集以量化隐藏的通信。
  3. 不要将内核耗时相加并称之为墙钟时间。并发内核可能由于SM或带宽争用而运行更长，即使暴露的GPU活动并集和端到端步进时间下降。
  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模型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"：在NVL8和NVL72系统上均对拓扑敏感，当HybridEP已安装时候选

HybridEP加上普通EP重叠是当前16×H100 Qwen3 30B-A3B标准形状的实测赢家，而标准256×H100 Qwen3 235B配方使用标准alltoall加重叠。在目标容器中基准测试后端兼容性和吞吐量；GPU名称和EP度都不能单独决定赢家。

如果all-to-all路径在剖析中可见，将调度器调优与以下组合：

  • --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图

仅在剖析显示有意义的主机/启动间隙后才使用CUDA图。对于无丢弃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重叠实验隔离了从244.039到287.305 TFLOPS/GPU的上升；后来的规范配方在额外HybridEP调优后达到299.352。它们回答不同的问题。

陷阱

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

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

  3. FP8 MFU可能看起来误导性地低：在切换精度模式时比较绝对吞吐量和MFU。

  4. CUDA图和重新计算交互：TE范围图通常与选择性重新计算配对，而非全面全重新计算。

  5. 大规模下并行折叠不是可选项：一旦注意力和专家层明显想要不同布局，单一共享TP或EP计划成为两者的负担。

  6. 求和内核时间不是暴露时间：验证重叠时使用区间并集和通信/计算交集。

  7. 仅基准语义不是生产验收：强制路由、合成数据或禁用优化器/检查点路径必须与训练等价结果分开披露和验证。

  8. 功能激活需要证据：当后端可能回退、图可能捕获而无帮助或较低精度配方可能错过预期内核时，配置转储不足。

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