NVIDIA NeMo AutoModel 借助专家并行和 DeepEP 将 MoE 微调速度提升 3.7 倍

NVIDIA 推出了 NeMo AutoModel，这是 NeMo 框架中的一个开源库，为混合专家（MoE）模型的微调带来了显著的性能提升。该库直接构建于 Hugging Face Transformers v5 之上，与最佳的 Transformers v5 配置相比，训练吞吐量提升了 3.4 至 3.7 倍，同时 GPU 内存消耗降低了 29% 至 32%，并保持了相同的 from_pretrained() API。用户只需更改一行导入代码即可获得这些好处。

背景：MoE 效率挑战

MoE 模型的兴起为高效训练带来了新的障碍。将 token 路由到数百个专家、将专家矩阵乘法融合到单个内核中、跨 GPU 分片权重以及将通信与计算重叠，这些都需要通用库通常不直接提供的基础设施。

Transformers v5 增加了对 MoE 的一流支持，包括专家后端、动态权重加载以及用于分布式执行的张量并行方案。但关键的性能部分，例如高效的通信重叠，仍未解决。NeMo AutoModel 填补了这一空白。

NeMo AutoModel 的工作原理

NeMo AutoModel 继承自 AutoModelForCausalLM，并加入了三项主要优化：专家并行（EP）、DeepEP 融合全到全分发和 TransformerEngine 内核。DeepEP 将通信与专家计算重叠，这是 Transformers v5 所缺乏的能力。由于 NeMo AutoModel 利用 v5 的可逆权重转换来加载模型，它将工程精力集中在可复用的核心操作上，而不是针对每个模型检查点的配置细节。

该库与 Hugging Face Transformers API 保持完全兼容。如 NVIDIA 的代码示例所示，加载模型只需更改 import 语句即可。

性能基准测试

多节点：Nemotron 3 Ultra 550B A55B

对于在 16 个 H100 节点（128 个 GPU）上对 550B 参数 Nemotron 3 Ultra A55B 模型进行全参数微调，将专家并行度设为 64 的 NeMo AutoModel 实现了每个 GPU 每秒 815 个 token、每个 GPU 每秒 293 TFLOPs 的性能，峰值内存为 58.2 GiB。关键的是，Transformers v5 在此规模下会内存溢出，因此在没有 EP 的情况下无法进行直接比较。

单节点：Qwen3-30B-A3B

在配备 8 块 H100 80GB GPU 的单节点上，NeMo AutoModel（EP=8）记录为每个 GPU 每秒 11,340 个 token，相较于 Transformers v5 的每个 GPU 每秒 3,075 个 token，提升了 3.69 倍。峰值内存从 68.2 GiB 降至 48.1 GiB，降幅达 29%。值得注意的是，由于不兼容的 FSDP 打包单个专家模块，Transformers v4 在此模型上完全死锁。

单节点：Nemotron 3 Nano 30B A3B

对于 30B Nemotron Nano 模型，NeMo AutoModel 达到了每个 GPU 每秒 15,421 个 token，相较于 v5 的每个 GPU 每秒 4,583 个 token，提升了 3.36 倍。峰值内存消耗从 62.1 GiB 降至 42.5 GiB，降幅达 32%。

加速来源

性能提升来自三项协同优化：

• 专家并行 将专家权重分片到多个 GPU 上。在 8 个 GPU 上设置 EP=8 时，每个 GPU 仅保留专家参数的八分之一。对于一个专家权重约为 55 GiB 的模型，EP 将每个 GPU 的占用空间缩小到约 6.8 GiB，从而能够训练仅使用 FSDP 方法会内存溢出的模型。
• DeepEP 将 token 分发和合并融合到优化的 GPU 内核中，使用通信与专家计算重叠，而不是使用单独的 AllGather 和 ReduceScatter 集合操作。
• TransformerEngine 内核 加速了融合注意力、线性层和 RMSNorm 等核心操作，为所有层类型提供一致的加速效果。

利用 Transformers v5 特性

NeMo AutoModel 构建在 v5 的多个关键特性之上，包括专家后端系统。v5 引入了 experts_implementation 参数，具有三个选项：eager（用于调试的 for 循环）、batched_mm（单次批处理 GEMM）和 grouped_mm（排序 token 执行）。NeMo AutoModel 通过使用 DeepEP 融合全到全分发，结合分组 GEMM 内核和 TransformerEngine 线性层，扩展了该功能。

通过 v5 的 WeightConverter 和 WeightRenaming 实现的动态权重加载，使得 MoE 检查点可以以融合的 3D 张量形式存储。超过 20 种模型类型使用了此机制，包括 Mixtral、Qwen2 MoE、Qwen3 MoE、DeepSeek V2/V3 和 OLMoE。这些转换是完全可逆的：save_pretrained() 生成标准 Hugging Face 格式的 safetensors，下游工具如 vLLM 和 SGLang 可以加载这些文件。

结论与可用性

NVIDIA NeMo AutoModel 为扩展模型训练的 Hugging Face 用户提供了一条零摩擦的升级路径。通过更改一行导入代码，用户即可获得一个速度快三倍以上且内存占用显著减少的模型实例。代码、配置和基准测试脚本可在 NeMo AutoModel 仓库中获取。此项工作的核心贡献者包括 Adil Asif、Hemil Desai、Alexandros Koumparoulis 和 Huiying Li。