使用NVIDIA NeMo AutoModel加速Transformers微调
Accelerating Transformers Fine-Tuning with NVIDIA NeMo AutoModel
NVIDIA NeMo AutoModel 基于 HuggingFace Transformers v5 构建，通过专家并行（EP）、DeepEP 融合 all-to-all 调度和 TransformerEngine 内核，在微调 MoE 模型时实现训练吞吐量提升 3.4-3.7 倍、GPU 内存减少 29-32%。在 16 个 H100 节点上对 Nemotron 3 Ultra 550B A55B 进行全参数微调时，TPS/GPU 达 815，峰值内存 58.2 GiB；单节点 8 个 H100 上，Qwen3-30B-A3B 和 Nemotron 3 Nano 30B-A3B 的 TPS/GPU 分别为 11,340 和 15,421。用户仅需更改一行导入代码即可使用相同 API。
](https://huggingface.co/adil-asif)
HuggingFace Transformers 已成为开源 AI 生态的基石，而最近的 Transformers v5 版本通过为混合专家（MoE）模型提供一等支持进一步强化了它，MoE 如今是前沿模型的主流架构。v5 提供了 MoE 的基础组件：专家后端、动态权重加载和分布式执行，使 MoE 可扩展且易于构建。
NVIDIA NeMo AutoModel 是 NVIDIA NeMo 框架的一部分，是一个用于大规模构建自定义生成式 AI 模型的开源库。NeMo AutoModel 干净地构建在 v5 之上，增加了专家并行（Expert Parallelism）、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，并添加专家并行（EP）、DeepEP 融合 all-to-all 调度和 TransformerEngine 内核，构建在 v5 之上。DeepEP 是 v5 尚未具备的部分：它将通信与专家计算重叠。并且由于 NeMo AutoModel 利用 v5 的可逆权重转换来加载每个模型，它可以将工程工作集中在这些可重用的核心操作上，而不是每个模型的 checkpoint 管道，而 save_pretrained() 仍然输出标准的 HF checkpoint，vLLM 和 SGLang 等工具可以加载这些 checkpoint。
下一节将介绍两者如何协同工作，以及我们测量的性能提升，范围从跨 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 attention、融合线性层和自定义专家内核。对于其他所有模型，它会回退到原生 HF，同时仍然应用诸如 Liger kernel 补丁等优化。无论采用哪种路径，生成的模型都准备好进行扩展：传入 device_mesh 即可进行多 GPU 训练，无需进一步重写。
NeMo AutoModel 真正闪耀之处在于将 MoE 模型扩展到多 GPU 训练。要使用专家并行在 8 个 GPU 上训练 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）。
方法：
| 参数 | 值 |
|---|---|
| 硬件 | 16x H100 80GB（128 个 GPU） |
| 专家并行 | EP=64 |
| 本地 batch size | 2 |
| 序列长度 | 4,096 |
| 特性 | MTP、activation checkpointing、融合线性交叉熵 |
| 内核 | DeepEP 调度 + torch_mm 专家 + TransformerEngine |
| 指标 | NeMo AutoModel (EP=64) |
|---|---|
| TPS/GPU（平均） | 815 |
| TFLOP/s/GPU | ~293 |
| 峰值内存 | 58.2 GiB |
为什么没有 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 + 自定义内核）。
方法：
| 参数 | 值 |
|---|---|
| 硬件 | 8x H100 80GB（单节点） |
| 序列长度 | 4,096 |
| 本地 batch size | 1 |
关于路由门的说明。 下面的 NeMo AutoModel 数据使用了平衡路由门，这迫使 token 均匀分布在各个专家之间。这模拟了 MoE 训练所追求的_理想_工作点：一个训练良好的模型的负载均衡损失会使专家利用率接近均匀，因此平衡路由反映了真实工作负载收敛到的稳态（并消除了随机虚拟 token 可能注入到专家并行中的掉队者噪声）。v4/v5 在相同的虚拟 token 上运行其原生路由器。因此，平衡门衡量的是 NeMo AutoModel 在其目标 MoE 工作点的表现，而 v4/v5 列反映的是它们开箱即用的行为。
Qwen3-30B-A3B
| 指标 | v4 | v5 (FA2 + grouped_mm) | NeMo AutoModel (EP=8) | v5 → NeMo AutoModel |
|---|---|---|---|---|
| TPS/GPU（平均） | 死锁 | 3,075 | 11,340 | 3.69x |
| 峰值内存 | — | 68.2 GiB | 48.1 GiB | -29% |
| 平均前向+损失 | — | 582 ms | 194 ms | 3.00x |
| 平均反向 | — | 758 ms | 178 ms | 4.26x |
为什么 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（hub 代码） | v5 (FA2 + grouped_mm + Mamba CUDA) | NeMo AutoModel (EP=8) | v5 → NeMo AutoModel |
|---|---|---|---|---|
| TPS/GPU（平均） | 1,807 | 4,583 | 15,421 | 3.36x |
| 峰值内存 | 61.9 GiB | 62.1 GiB | 42.5 GiB | -32% |
| 平均前向+损失 | 1,024 ms | 283 ms | 109 ms | 2.60x |
| 平均反向 | 1,246 ms | 611 ms | 157 ms | 3.89x |
v4 配置： trust_remote_code=True（NVIDIA 的 hub 建模代码）。hub 代码的专家循环是 FSDP 安全的（无论 token 分配如何，都迭代所有专家），因此它不会像 Qwen3 v4 那样死锁。
加速的来源
NeMo AutoModel 相对于 Transformers v5 的 3.4-3.7 倍加速来自三个方面：
专家并行减少了内存压力。 EP=8 将专家权重分布到多个 GPU 上，将每个 GPU 的 MoE 内存占用减少 8 倍。对于 Qwen3，这使峰值内存从 68.2 GiB 降至 48.1 GiB（-29%）。对于 Nemotron Nano，从 62.1 GiB 降至 42.5 GiB（-32%），为更大的 batch size 或更长的序列腾出了空间。
DeepEP 将通信与计算融合。 DeepEP 不是为专家路由使用单独的 AllGather/ReduceScatter 集合，而是将 token 调度和组合融合到优化的 GPU 内核中，使通信与专家计算重叠。
TransformerEngine 内核加速了核心操作。 TE 的融合 attention、线性层和 RMSNorm 实现，在所有层类型（不仅仅是 MoE 层）上，都提供了相对于其 PyTorch/Flash Attention 对应物的持续加速。
HuggingFace AutoModel 利用的 Transformers v5 特性
专家后端
Transformers v5 中最有影响力的特性之一是 experts_implementation 参数，它包含三个专家后端：
| 后端 | 描述 | 最佳用途 |
|---|---|---|
| eager | 对选定的专家进行 for 循环 | 调试、兼容性和正确性。v4 也可用。 |
| batched_mm | 复制专家参数，通过 torch.bmm 进行单个批处理 GEMM | 小输入，使用 torch.compile 时速度快。v5 新增。 |
| grouped_mm | 按专家对 token 排序，通过 torch.nn.functional.grouped_mm 进行单个分组 GEMM | 训练（内存高效，无参数复制）。v5 新增。 |
grouped_mm 后端是关键的训练优化：它不是逐个循环专家，而是按分配的专家对 token 排序，并执行单个融合的分组矩阵乘法。
NeMo AutoModel 更进一步。对于具有自定义实现的模型，它使用 DeepEP 融合 all-to-all 调度，结合分组 GEMM 内核和 TransformerEngine 线性层。演进过程如下：
v4 (eager for 循环) → 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 线性层
experts="torch_mm", # 分组专家矩阵乘法
dispatcher="deepep", # DeepEP 融合 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，与数据并行网格并列（而不是从中划分），使用 PyTorch 的 DTensor 和 Shard(0)。由于专家网格与数据并行是正交的，两者可以在同一设备上组合。在 8 个 GPU 上，NeMo AutoModel 同时运行 ep=8 和 dp=8，因此每个 GPU 都在自己的数据分片上训练，同时只持有 1/8 的专家。专家权重沿专家维度物理分片到多个 GPU 上。
# 来自 nemo_automodel/components/moe/parallelizer.py
from torch.distributed.tensor import Shard, distribute_tensor
# 每个 GPU 只持有 1/ep_size 的专家权重
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 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
结论
NVIDIA NeMo AutoModel 是 HuggingFace 用户扩展模型训练规模的自然下一步。通过直接构建在 Transformers v5 之上，AutoModel 提供了一个零摩擦的升级路径：更改一行导入代码，即可获得一个速度提升超过三倍的模型实例。
在 Qwen3-30B-A3B 和 Nemotron 3 Nano 30B-A3B 上，与最佳的 Transformers v5 配置相比，这带来了 3.4-3.7 倍的训练吞吐量提升和 29-32% 的 GPU 内存减少。并且由于真正的专家并行将专家分片到多个 GPU 上，相同的路径可以扩展到跨 16 个节点全参数微调像 Nemotron 3 Ultra 这样的 550B 模型，这是专家并行对于将模型装入内存至关重要的场景。由于 NeMo AutoModel checkpoint 是标准的 HF 格式 safetensors，您可以将它们部署在 vLLM 和 SGLang 等推理框架上。
代码、配置和基准测试脚本均可在 NeMo AutoModel 仓库中找到。
致谢
此项工作的核心贡献者，按姓氏字母顺序排列：Adil Asif、Hemil Desai、Alexandros Koumparoulis 和 Huiying Li。