使用 NVIDIA NeMo AutoModel 加速 Transformer 架构微调

HuggingFace Transformers has become the foundation of the open-source AI ecosystem, and the recent Transformers v5 release strengthened it with first-class support for Mixture-of-Experts (MoE) models, now the dominant architecture for frontier models. v5 ships the MoE foundations: expert backends, dynamic weight loading, and distributed execution that make MoE extensible and easy to build on.

NVIDIA NeMo AutoModel 是 NVIDIA NeMo 框架中的一个开源库，用于大规模构建定制化生成式 AI 模型。NeMo AutoModel 直接构建在 Transformers v5 之上，新增了专家并行（Expert Parallelism）、DeepEP 融合全对全（all-to-all）调度以及 TransformerEngine 内核，并利用 v5 的动态权重加载，将这些优化推广到越来越广泛的模型族。其成果是：在对 MoE 模型进行微调时，训练吞吐量提升 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`，并在其上添加专家并行（EP）、DeepEP 融合全对全调度以及 TransformerEngine 内核，构建于 v5 之上。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注意力机制、融合线性层和自定义专家内核。对于其他架构，它会回退到标准HF实现，同时仍然应用Liger内核补丁等优化手段。无论走哪条路径，生成的模型都具备扩展能力：只需传入 device_mesh 参数，即可实现多GPU训练，无需额外重写代码。

NeMo AutoModel的真正优势在于将MoE模型扩展到多GPU训练。要在8张GPU上采用专家并行（Expert Parallelism）训练Nemotron 3 Nano 30B A3B，只需添加分布式网格配置：

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()

通过一次 from_pretrained() 调用，即可获得由FSDP2、专家并行、TransformerEngine内核和DeepEP调度带来的速度、可扩展性和内存优化。

性能对比

我们在两种场景下评估了NeMo AutoModel：在16个节点上对前沿规模的550B模型进行全量微调，以及在单个节点上训练两个30B MoE模型。550B的结果说明了专家并行在大规模训练中的必要性；30B的结果则量化了每张GPU相对于Transformers v5的加速效果。

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 基准测试

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

方法论：

关于路由门的说明。下方的 NeMo AutoModel 数据使用了均衡路由门，这迫使 token 在各专家之间均匀分布。这模拟了 MoE 在训练中趋向的理想操作点：一个训练良好的模型的负载均衡损失会使专家利用率趋近均匀，因此均衡路由反映了真实工作负载收敛到的稳态（并消除了随机虚拟 token 原本会注入专家并行中的拖尾噪声）。v4/v5 对同样的虚拟 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. 专家并行（Expert Parallelism）降低了内存压力。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 内核，使通信与专家计算重叠进行。

3. TransformerEngine 内核加速核心运算。TE 的融合注意力、线性层和 RMSNorm 实现，在所有层类型（不仅仅是 MoE 层）上相比其 PyTorch / Flash Attention 等效实现都提供了一致的加速效果。

HuggingFace AutoModel 利用的 Transformers v5 特性

专家后端（Expert Backends）

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（独立于数据并行网格，而非从中划分），并使用 PyTorch 的 DTensor 配合 Shard(0)。由于专家网格与数据并行正交，两者可在同一设备上组合。在 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 内核中，并在与分组 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 上，与最佳 Transformer v5 配置相比，该方案可实现 3.4-3.7 倍的训练吞吐量提升，同时减少 29-32% 的 GPU 显存占用。由于真正的专家并行（Expert Parallelism）会将专家分片到多个 GPU 上，相同的路径可以扩展到在 16 个节点上对像 Nemotron 3 Ultra 这样的 550B 模型进行全参数微调，这正是专家并行对于将模型装入显存至关重要的场景。因为 NeMo AutoModel 检查点是标准 HF 格式的 safetensors，你可以将它们部署在 vLLM 和 SGLang 等推理框架上。

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

致谢

本研究的主要贡献者（按姓氏字母顺序排列）：Adil Asif、Hemil Desai、Alexandros Koumparoulis 和 Huiying Li。

本文提及的模型 2