NVIDIA NeMo AutoModel：一行代码加速Transformer 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：只需一行 import，无需其他代码改动。

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

背景

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

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

NeMo AutoModel 构建在 v5 之上，通过继承 AutoModelForCausalLM，并加入专家并行（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。

下面是在两者中加载模型的样子。只有导入语句不同：

这一个导入语句做了很多工作。对于像 Qwen3、NVIDIA Nemotron、GPT-OSS 和 DeepSeek V3 这样的热门 MoE 架构，NeMo AutoModel 提供了手工调优的实现，包含 TransformerEngine 注意力、融合线性层和自定义专家 kernel。对于其他所有情况，它会回退到原版 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 kernel 和 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 的专家并行（Expert Parallelism）将专家分片到多块 GPU 上，使内存占用控制在预算之内，这正是让完整微调得以运行的原因。下面的 30B 对比展示了在 v5 能够运行的场景下同样的优势。

单节点 30B MoE 基准测试

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

方法说明：

关于路由门控的说明。下面的 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.7x 加速来自三个方面：

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 将通信与计算融合。DeepEP 不再为专家路由使用单独的 AllGather/ReduceScatter 集合通信，而是将 token 分发与合并融合进优化后的 GPU kernel，使通信与专家计算相互重叠。

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

HuggingFace AutoModel 所利用的 Transformers v5 特性

专家后端

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

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

NeMo AutoModel 更进一步。对于带有自定义实现的模型，它使用 DeepEP 融合 all-to-all 分发，结合分组 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",           
    linear="te",         
    experts="torch_mm",  
    dispatcher="deepep", 
)

专家并行与 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 torch.distributed.tensor import Shard, distribute_tensor


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 kernel，并与分组 GEMM 结合用于分组专家计算，从而带来显著的加速。在我们的 大规模 MoE 基准测试中，与 all-gather + 循环专家基线相比，DeepEP + 分组 GEMM 在完整的 DeepSeek V3 671B 模型上将每次迭代的成本降低了 47%。

动态权重加载

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

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

快速上手

如需试用 NeMo AutoModel，请访问我们的官方文档页面。

更多详情，请参阅：

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

结论

对于希望扩大模型训练规模的 HuggingFace 用户来说，NVIDIA NeMo AutoModel 是顺理成章的下一步。AutoModel 直接构建在 Transformers v5 之上，提供了一条零阻力的升级路径：只需改动一行 import，就能获得一个速度提升三倍以上的模型实例。

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

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

致谢

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