⚡

长上下文MoE训练优化Skill nemo-mbridge-perf-moe-long-context

该技能提供MoE模型在长序列（如128K以上）训练时的优化指南，涵盖上下文并行（CP）大小调整、选择性重计算、调度器选择，以及基于DSV3、Qwen3、Qwen3-Next等实际实验的配置模式，帮助解决长上下文训练中的OOM、吞吐下降等问题。适用于Megatron Bridge框架。关键词：长上下文MoE、CP大小、选择性重计算、Megatron Bridge、128K、DeepEP、HybridEP。

Megatron-Core训练 0 次安装 1 次浏览 更新于 9/7/2026
名称 nemo-mbridge-perf-moe-long-context
描述 Megatron Bridge 的长上下文MoE训练指导。涵盖CP大小调整、选择性重计算、调度器选择，以及来自DSV3、Qwen3和Qwen3-Next长上下文实验的实践经验。
开源协议 Apache-2.0 when_to_use: 在长序列长度下训练MoE，或调查导致长上下文MoE OOM或吞吐下降的提交；‘long context MoE’、‘128k tokens’、‘CP sizing for long sequences’、‘selective recompute long context’、‘MoE long-context OOM’。

MoE 长上下文训练

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

长上下文的变化

一旦序列长度远超4K类范围，注意力内存和激活驻留成为主导约束。对于MoE模型，这通常意味着你需要以下某些组合：

  • 上下文并行（context parallelism）
  • 选择性重计算（selective recompute）
  • 较低精度
  • 优化器状态的CPU卸载
  • 不会浪费剩余较小DP预算的调度器和流水线并行布局

实测扩展模式

DSV3在H100上

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

  • 一旦超过最短上下文，选择性重计算优于全重计算
  • 如果适当增加CP，吞吐量在中长到很长上下文中保持在较窄的范围内
  • 随着CP增长，权衡从“内存适配”转向“GPU数量可行性”

换句话说，只要布局选择得当，长上下文不会立即导致利用率下降，但它确实非常快地消耗了DP预算。

Qwen3-Next在GB200上

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

  • 8K和32K在适度CP下仍然实用
  • 64K可行，但吞吐量下降明显，内存变得更加紧张
  • 流水线布局和分组GEMM改进几乎与CP一样重要

Qwen3 235B在GB200上

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-heavy重计算：重计算注意力内部可能增加大量工作，但相比重计算较小的MoE和MLP侧模块，内存收益却更少。

  5. 在NVL72系统上把TP作为另一个杠杆：GB200和GB300运行有时可以用一些CP换取TP，同时仍保持高效。

  6. 假设GBS需要缩小：随着CP增加和DP减少，你可能需要减少全局批量大小或接受更高的梯度累积。

代表性配置族

DSV3 128K H100

TP=1  CP=32  EP=32  PP=8  VPP=4
精度: FP8类
调度器: DeepEP
重计算: up_proj, norm, moe, mlp
额外内存帮助: 优化器CPU卸载

DSV3 256K H100

TP=1  CP=64  EP=32  PP=8  EDP=2  VPP=4
精度: FP8类
调度器: DeepEP
重计算: up_proj, norm, moe, mlp
额外内存帮助: 优化器CPU卸载

Qwen3 235B 128K GB200

TP=4  CP=4  EP=32  PP=4  VPP=12
精度: BF16或MXFP8
调度器: HybridEP
重计算: moe_act, norm
CUDA图: attn + moe_router + moe_preprocess

重计算与CUDA图指导

对于长上下文MoE训练：

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

有用的参考：

  • @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图需要静态形状：可变长度批次和机会性填充策略可能会悄悄破坏路径。

  5. 容器和内核支持在128K+时更为重要：长上下文路径往往比短上下文启动时依赖更新的内核和错误修复。