mcpField
N
nemo-mbridge-perf-moe-long-context

nemo-mbridge-perf-moe-long-context

github.com/nvidia开发3.8k1.8k 次安装今天更新
代码仓库

关于 nemo-mbridge-perf-moe-long-context

面向训练工程师的长上下文 MoE 训练调优技能，说明并行策略、重计算、精度与卸载之间的取舍，并提供多种模型和硬件配置示例。

主要功能

  1. 01上下文并行度 sizing 指引
  2. 02选择性重计算建议
  3. 03精度与优化器卸载选项
  4. 04H100 和 GB200 配置示例
  5. 05CUDA Graph 使用指南

适用场景

  1. 01规划长上下文 MoE 训练
  2. 02选择 CP、TP、EP 和 PP 设置
  3. 03降低激活或优化器内存占用
  4. 04评估 GPU 数量可行性

安装

复制下面的命令，在终端中运行。

$npx skills add https://github.com/nvidia/skills --skill nemo-mbridge-perf-moe-long-context

安装后在 Claude Code 中输入/nemo-mbridge-perf-moe-long-context确认已加载。

包含的文件

除 SKILL.md 外，安装时一起下载的参考资料和脚本。

BENCHMARK.md
card.yaml
skill-card.md
skill.oms.sig
evals/
evals.json

详细介绍

MoE 长上下文训练

稳定版文档：@docs/training/moe-optimization.md 说明卡片：@skills/nemo-mbridge-perf-moe-long-context/card.yaml

长上下文会带来哪些变化

当序列长度远超 4K 级别后，注意力内存和激活驻留会成为主要瓶颈。对于 MoE 模型，这通常意味着需要组合使用以下方法：

  • 上下文并行
  • 选择性重计算
  • 更低精度
  • 将优化器状态卸载到 CPU
  • 采用不会浪费所剩不多 DP 预算的分发器和 PP 布局

大致的扩展规律

H100 上的 DSV3

DSV3 长上下文运行显示出稳定的规律：

  • 超过最短上下文后，选择性重计算的效果优于完整重计算
  • 只要适当提高 CP，从中等长度到超长上下文，吞吐量都能维持在较窄的范围内
  • 随着 CP 增大，权衡重点会从“内存是否足够”转向“GPU 数量是否可行”

换言之，只要布局选择得当，长上下文不会立即导致利用率骤降，但会很快耗尽 DP 预算。

GB200 上的 Qwen3-Next

Qwen3-Next 更像是一个对内存敏感的中等规模模型：

  • 8K 和 32K 在适度 CP 下仍然实用
  • 64K 可以实现，但吞吐量下降明显，内存也会变得非常紧张
  • 流水线布局和分组 GEMM 优化几乎与 CP 同样重要
GB200 上的 Qwen3 235B

Qwen3 235B 表明，在 NVL72 系统上，只要协调好 TP、CP 和 HybridEP，长上下文仍然可以高效运行。最佳的 128K 级配置不只是“能跑起来”的方案；只要平衡好路由、并行方式和重计算，也能保持很高的效率。

CP 设置经验法则

  1. 以 4K 分片为目标起步：一个不错的初始估算是 CP ~= seq_len / 4096，然后调整为实际可用的 2 的幂布局。

  2. 尽可能保留 DP：当 CP、EP、TP 和 PP 合起来将 DP 挤压到最低限度时，长上下文扩展会变得脆弱。

  3. 优先选择性重计算：在采用完整重计算前，先对 up_proj、norm、moe、moe_act 或 mlp 等模块进行重计算。

  4. 极长上下文下避免大量使用 SDPA 的重计算：重计算注意力内部操作所增加的计算量，可能远超重计算较小的 MoE 和 MLP 侧模块所节省的内存。

  5. 在 NVL72 系统上将 TP 作为另一个调节手段：GB200 和 GB300 的运行有时可以用部分 TP 替代 CP，同时保持较高效率。

  6. 假设 GBS 需要缩小：随着 CP 增大、DP 减少，可能需要降低全局批次大小，或接受更高的 GA。

代表性配置系列

H100 上 128K 的 DSV3
TP=1  CP=32  EP=32  PP=8  VPP=4
Precision: FP8-class
Dispatcher: DeepEP
Recompute: up_proj, norm, moe, mlp
Extra memory help: optimizer CPU offload
H100 上 256K 的 DSV3
TP=1  CP=64  EP=32  PP=8  EDP=2  VPP=4
Precision: FP8-class
Dispatcher: DeepEP
Recompute: up_proj, norm, moe, mlp
Extra memory help: optimizer CPU offload
GB200 上 128K 的 Qwen3 235B
TP=4  CP=4  EP=32  PP=4  VPP=12
Precision: BF16 or MXFP8
Dispatcher: HybridEP
Recompute: moe_act, norm
CUDA Graph: attn + moe_router + moe_preprocess

重计算和 CUDA Graph 指南

对于长上下文 MoE 训练：

  • 从选择性重计算开始
  • 只有在形状和路由路径稳定后，才添加 CUDA Graph
  • 使用 CUDA Graph 时，保持序列长度和 MBS 不变
  • 如果运行依赖高度动态的批次，优先使用即时执行

有用的参考资料：

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

注意事项

  1. CP 不能替代 EP 或 PP：它会增加一个维度，不会让其他维度消失。

  2. 良好的 4K 基线仍可能不适用于长上下文：路由模式、重计算选择和卸载策略通常都需要调整。

  3. GPU 数量的可行性会成为真正的限制：单看某个配置时，极长上下文似乎可行；但在整个模型中如实加入 EP 和 PP 后，可能就无法实现。

  4. CUDA Graph 需要静态形状：可变长度批次和临时填充策略可能会在不显眼的情况下破坏执行路径。

  5. 在 128K 及以上场景中，容器和内核支持更加重要：与短上下文启动相比，长上下文路径往往依赖更新的内核和错误修复。

评论

支持 Markdown