NVIDIA NeMo AutoModel 发布，MoE 微调吞吐提升 3.4-3.7 倍

HuggingFace Transformers 已成为开源 AI 生态系统的基石，而最近的 Transformers v5 版本通过为混合专家（MoE）模型提供一流支持进一步强化了这一点，MoE 现已成为 前沿模型 的主流架构。v5 提供了 MoE 基础：专家后端、动态权重加载和分布式执行，使 MoE 可扩展且易于构建。

NVIDIA NeMo AutoModel 是 NVIDIA NeMo 框架 中的一个开源库，用于大规模构建自定义生成式 AI 模型。NeMo AutoModel 在 v5 之上干净地构建，增加了专家并行、DeepEP 融合 all-to-all 调度和 TransformerEngine 内核，并依赖 v5 的动态权重加载将这些优化带给广泛且不断增长的模型系列。其收益是，在微调 MoE 模型时，相比原生 Transformers v5，训练吞吐量提高 3.4-3.7 倍，GPU 内存减少 29-32%，且使用相同的 from_pretrained() API：只需一行导入，无需其他代码更改。

本博客详细介绍了这种组合如何工作，以及用户如何在不更改 API 的情况下更快地微调 MoE 模型。

背景

MoE 模型的兴起为高效训练带来了新的挑战：在数百个专家之间路由 token、将专家矩阵乘法融合到单个内核中、跨 GPU 分片权重，以及重叠通信与计算，这些都需要通用库开箱即用之外的基础设施。

Transformers v5（“v5”）引入了一流的 MoE 支持，例如 专家后端、动态权重加载 以及用于分布式执行的张量并行计划。此外，v5 通过将 PyTorch 的 DeviceMesh 直接集成到 from_pretrained() 中，使分布式训练成为一等公民。

NeMo AutoModel 通过继承 AutoModelForCausalLM 在 v5 之上构建，并增加了专家并行（EP）、DeepEP 融合 all-to-all 调度和 TransformerEngine 内核。DeepEP 是 v5 尚未具备的部分：它将通信与专家计算重叠。而且，由于 NeMo AutoModel 利用 v5 的可逆权重转换来加载每个模型，它可以将工程重点放在这些可复用的核心操作上，而不是每个模型的检查点管道，同时 save_pretrained() 仍会输出标准 HF 检查点，供 vLLM 和 SGLang 等工具加载。

下一节将介绍两者如何协同工作以及我们测得的性能提升，从跨 16 个节点全量微调 NVIDIA Nemotron 3 Ultra 550B A55B 到单节点模型如 Qwen3-30B-A3B 和 Nemotron 3 Nano 30B A3B。

NeMo AutoModel：相同 API，更高性能

NeMo AutoModel 的目标之一是与 HuggingFace Transformers 的 API 兼容，以支持开源社区。NeMoAutoModelForCausalLM 继承自 AutoModelForCausalLM，因此任何适用于 HF 模型的代码也适用于 AutoModel。

以下是两者加载模型的样子。只有导入发生了变化：

这一行 import 就完成了大量工作。对于 Qwen3、NVIDIA Nemotron、GPT-OSS 和 DeepSeek V3 等流行的 MoE 架构，NeMo AutoModel 提供了手工调优的实现，包含 TransformerEngine 注意力、融合线性层和自定义专家内核。对于其他所有情况，它会回退到原生 HF，同时仍然应用诸如 Liger kernel 补丁等优化。无论走哪条路径，得到的模型都已准备好扩展：传入 device_mesh，即可进行多 GPU 训练，无需进一步改写。

NeMo AutoModel 真正出彩的地方在于将 MoE 模型扩展到多 GPU 训练。要在 8 个 GPU 上以专家并行训练 Nemotron 3 Nano 30B A3B，只需添加分布式 mesh 配置：

import os
import torch
import torch.distributed as dist
from nemo_automodel import NeMoAutoModelForCausalLM
from nemo_automodel.recipes._dist_utils import create_distributed_setup_from_config

dist.init_process_group(backend="nccl")
torch.manual_seed(0)
torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))

dist_setup = create_distributed_setup_from_config(
    {
        "strategy": "fsdp2",
        "ep_size": 8,
    },
)

model = NeMoAutoModelForCausalLM.from_pretrained(
    "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
    dtype=torch.bfloat16,
    distributed_setup=dist_setup,
)

dist.destroy_process_group()

这通过 FSDP2、专家并行、TransformerEngine 内核和 DeepEP 调度带来了速度、可扩展性和内存优化，全部来自一次 from_pretrained() 调用。

性能对比

我们在两种场景下评估了 NeMo AutoModel：在 16 个节点上全量微调一个前沿规模的 550B 模型，以及在单个节点上训练两个 30B MoE 模型。550B 的结果表明为什么专家并行在大规模下至关重要；30B 的结果量化了相对于 Transformers v5 的每 GPU 加速比。

Nemotron 3 Ultra 550B A55B（全量微调，多节点）

Nemotron 3 Ultra 550B A55B 是一个 550B 参数的混合模型，搭载 Mamba2、LatentMoE 和多 Token 预测（MTP）。我们基准测试了全量微调：每个参数都会更新，Adam 优化器状态也会物化，在这个规模下需要 16 个 H100 节点（128 个 GPU）。

方法：

为什么没有 Transformers v5 列。 Transformers v5 在这个规模下会内存不足，因此这里没有 v5 数字可报告。AutoModel 的专家并行将专家分片到多个 GPU 上，使占用空间控制在预算内，这正是全量微调能够运行的原因。下面的 30B 对比展示了 v5 能容纳时的同样优势。

单节点 30B MoE 基准测试

我们在单个节点上使用 8x H100 80GB GPU 对三种方法进行了基准测试：HF Transformers v4（hub 代码）、HF Transformers v5（采用可用的最佳优化）和 NeMo AutoModel（EP=8 + 自定义内核）。

方法：

关于路由门控的说明。 下面的 NeMo AutoModel 数字使用平衡路由门控，这会强制 token 在专家之间均匀分布。这模拟了 MoE 训练所趋向的理想工作点：训练良好的模型的负载均衡损失会将专家利用率推向接近均匀，因此平衡路由反映了真实工作负载收敛到的稳态（并消除了随机 dummy token 否则会注入专家并行的拖尾噪声）。v4/v5 在相同的 dummy token 上运行其原生路由器。因此，平衡门控在 NeMo AutoModel 的目标 MoE 工作点上对其进行测量，而 v4/v5 列反映其开箱即用行为。

Qwen3-30B-A3B

为什么 v4 会死锁：Transformers v4 将 Qwen3 MoE 专家存储为 128 个独立 MLP 模块组成的 ModuleList，每个模块单独进行 FSDP 封装。前向传播使用数据依赖的循环，仅迭代接收到 token 的专家。由于每个 rank 的数据不同，不同 rank 会跳过不同的专家，导致 FSDP AllGather/ReduceScatter 集合通信不匹配，从而无限期挂起。Transformers v5 通过将专家存储为融合的 3D 参数张量（没有逐专家模块，没有逐专家 FSDP 集合通信）修复了此问题。

Nemotron 3 Nano 30B A3B

v4 配置：trust_remote_code=True（NVIDIA 的 hub 建模代码）。hub 代码的专家循环是 FSDP 安全的（无论 token 分配如何都会迭代所有专家），因此不会像 Qwen3 v4 那样死锁。

加速从何而来

NeMo AutoModel 相比 Transformers v5 带来的 3.4-3.7 倍加速来自三个方面：

1. 专家并行降低内存压力。EP=8 将专家权重分布到多个 GPU 上，将每 GPU 的 MoE 占用减少 8 倍。对于 Qwen3，这将峰值内存从 68.2 GiB 降至 48.1 GiB（-29%）。对于 Nemotron Nano，从 62.1 GiB 降至 42.5 GiB（-32%），为更大的批大小或更长的序列释放了空间。

2. DeepEP 将通信与计算融合。不再为专家路由使用单独的 AllGather/ReduceScatter 集合通信，DeepEP 将 token 分发和合并融合到优化的 GPU 内核中，使通信与专家计算重叠。

3. TransformerEngine 内核加速核心操作。TE 的融合注意力、线性层和 RMSNorm 实现在所有层类型上都比 PyTorch/Flash Attention 的对应实现提供一致的加速，而不仅仅是 MoE 层。

HuggingFace AutoModel 利用的 Transformers v5 特性

专家后端

Transformers v5 中最具影响力的特性之一是 experts_implementation 参数，它包含三种专家后端：

grouped_mm 后端是关键训练优化：不是逐个循环专家，而是按分配的专家对 token 排序，并执行单次融合分组矩阵乘法。

NeMo AutoModel 更进一步。对于具有自定义实现的模型，它使用 DeepEP 融合全对全分发，结合分组 GEMM 内核和 TransformerEngine 线性层。演进过程如下：

v4 (eager for-loop) → v5 (grouped_mm) → NeMo AutoModel (DeepEP + GMM + TE)

在 NeMo AutoModel 中，专家后端通过 BackendConfig 配置：

from nemo_automodel.components.models.common.utils import BackendConfig

backend = BackendConfig(
    attn="te",           # TransformerEngine attention
    linear="te",         # TransformerEngine linear layers
    experts="torch_mm",  # Grouped expert matmul
    dispatcher="deepep", # DeepEP fused all-to-all
)

专家并行与 DeepEP

Transformers v5 还提供了 专家并行路径。它将专家权重分片到多个 GPU 上。GroupedGemmParallel 风格只加载每个设备的本地专家，而 RouterParallel 负责路由 token 并通过 all_reduce 合并结果。它巧妙地构建在 v5 现有的张量并行机制之上。启用它后，模型的 tp_plan 会返回其 专家计划，因此专家并行与数据并行共享设备预算（ep × dp = world_size）。对于这里的单节点 30B 基准测试，我们发现纯数据并行的 v5（dp=8，ep=1）是最快的 v5 配置，因此这就是我们报告的 v5 设置。

NeMo AutoModel 采用了一种互补的方法，专为多 GPU MoE 训练而调优。它将 EP 作为自己的并行维度，一个专用的 moe_mesh 与数据并行 mesh 并列（而非从中划分），使用 PyTorch 的 DTensor 和 Shard(0)。由于专家 mesh 与数据并行正交，两者可以在相同的设备上组合。在 8 个 GPU 上，NeMo AutoModel 同时运行 ep=8 和 dp=8，因此每个 GPU 都在自己的数据分片上训练，同时只持有 1/8 的专家。专家权重沿专家维度在 GPU 之间进行物理分片。

# From nemo_automodel/components/moe/parallelizer.py
from torch.distributed.tensor import Shard, distribute_tensor

# Each GPU holds only 1/ep_size of the expert weights
distribute_tensor(param, device_mesh, [Shard(0)])

在 8 个 GPU 上使用 ep_size=8 时，每个 GPU 只持有 1/8 的专家参数。对于像 Nemotron-3-Nano-30B-A3B 这样拥有约 55 GiB 专家权重的模型，EP 将每 GPU 的专家占用从约 55 GiB 减少到约 6.8 GiB，使得在仅使用 FSDP 的方法内存不足的情况下也能进行训练。

在 EP 之上，NeMo AutoModel 集成了 DeepEP，它将 token 路由融合到优化的 GPU 内核中，并与分组 GEMM 结合用于分组专家计算时带来显著的加速。在我们的大规模 MoE 基准测试中，与 all-gather + 循环专家基线相比，DeepEP + 分组 GEMM 在完整的 DeepSeek V3 671B 模型上将每次迭代的成本降低了 47%。

动态权重加载

Transformers v5 还通过 WeightConverter 和 WeightRenaming 引入了动态权重加载系统。这使得 MoE 检查点可以存储为融合的 3D 张量，以实现更高效的执行。WeightConverter 在 from_pretrained() 期间应用可组合的操作来即时转换检查点张量。

NeMo AutoModel 是此 v5 API 的直接使用者。超过 20 种模型类型通过 MODELS_REQUIRING_TENSOR_MERGING 使用此机制，包括 Mixtral、Qwen2 MoE、Qwen3 MoE、DeepSeek V2/V3、OLMoE 等。这些转换完全可逆：save_pretrained() 生成标准 HF 格式的检查点，任何下游工具都可以加载。

快速开始

要试用 NeMo AutoModel，请访问我们的官方文档页面快速开始。

更多详情，请参阅：

• NeMo AutoModel HuggingFace API 兼容性指南
• NeMo AutoModel 模型覆盖范围
• NeMo AutoModel 性能摘要
• HuggingFace 上的 NeMo AutoModel

结论

NVIDIA NeMo AutoModel 是 HuggingFace 用户扩展模型训练的自然下一步。通过直接构建在 Transformers v5 之上，AutoModel 提供了零摩擦的升级路径：更改一行导入代码，即可获得速度提升三倍以上的模型实例。

在 Qwen3-30B-A3B 和 Nemotron 3 Nano 30B-A3B 上，与最佳的 Transformers v5 配置相比，这带来了 3.4-3.7 倍的训练吞吐量提升，同时 GPU 内存占用减少 29-32%。而且，由于真正的专家并行（Expert Parallelism）将专家分片到多个 GPU 上，同一条路径可以扩展到跨 16 个节点对 550B 模型（如 Nemotron 3 Ultra）进行全量微调，这正是专家并行成为将模型装入内存所必需的情形。由于 NeMo AutoModel 检查点是标准 HF 格式的 safetensors，你可以将它们部署到 vLLM 和 SGLang 等推理框架上。

代码、配置和基准测试脚本均可在 NeMo AutoModel 仓库中获取。

致谢

本工作的核心贡献者，按姓氏字母顺序排列：Adil Asif、Hemil Desai、Alexandros Koumparoulis 和 Huiying Li。