Accelerating Transformers Fine-Tuning with NVIDIA NeMo AutoModel

使用 NVIDIA NeMo AutoModel 加速 Transformer 微调

返回文章列表

[0

[-1

企业

+

]

文章

发布于 2026 年 6 月 24 日

点赞

26

[

• +20

Adil Asif

adil-asif

关注

nvidia

Alexandros Koumparoulis

akoumpa

Wenwen Gao

wgao2021

Sylendran Arunagiri

Sylendran95

David Messina

davidsalmessina70

Bernard Nguyen

bernardwin

HuggingFace Transformers 已经成为开源 AI 生态系统的基础，而最近发布的 Transformers v5 通过提供对 Mixture-of-Experts (MoE) 模型的一流支持，进一步加强了其地位，目前 MoE 是前沿模型的主要架构。v5 提供了 MoE 的基础：专家后端、动态权重加载和分布式执行，这些功能使 MoE 可扩展且易于构建。

NVIDIA NeMo AutoModel 是 NVIDIA NeMo 框架的一部分，是一个用于大规模构建自定义生成 AI 模型的开源库。NeMo AutoModel 在 v5 的基础上构建，添加了专家并行（Expert Parallelism）、DeepEP 融合的 all-to-all 分发和 TransformerEngine 内核，并依赖于 v5 的动态权重加载，将这些优化扩展到一个广泛且不断增长的模型家族。其优势在于，与原生 Transformers v5 相比，在微调 MoE 模型时，训练吞吐量提高了 3.4-3.7 倍，GPU 内存使用减少了 29-32%，并且使用相同的 from_pretrained() API：只需一行导入语句，无需其他代码更改。

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

背景

MoE 模型的兴起给高效训练带来了新的挑战：在数百个专家之间路由标记、将专家 matmuls 融合到单个内核、在 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 注意机制、融合的线性层和自定义专家内核。对于其他所有情况，它会退回到普通的 HF，同时仍然应用诸如 Liger 内核修补等优化。无论选择哪条路径，生成的模型都准备好进行扩展：只需传递一个 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()

这通过 FSDP2、专家并行、TransformerEngine 内核和 DeepEP 分发提供了速度、可扩展性和内存优化，所有这些都来自一个 from_pretrained() 调用。

性能比较

我们在两种情况下评估了 NeMo AutoModel：在 16 个节点上对一个前沿规模的 550B 模型进行完整微调，以及在一个节点上训练两个 30B MoE 模型。550B 的结果展示了为什么在大规模情况下专家并行是必不可少的；30B 的结果则量化了每个 GPU 相对于 Transformers v5 的加速效果。

Nemotron 3 Ultra 550B A55B（完整微调，多节点）

Nemotron 3 Ultra 550B A55B 是一个 550B 参数的混合模型，包含 Mamba2、LatentMoE 和多令牌预测（MTP）。我们对完整微调进行了基准测试：每个参数都会被更新，Adam 优化器状态会被实例化，这在该规模下涉及 16 个 H100 节点（128 个 GPU）。

方法：

| 参数 | 值 | |------|----| | 硬件 | 16x H100 80GB（128 个 GPU） | | 专家并行 | EP=64 | | 本地批量大小 | 2 | | 序列长度 | 4,096 | | 功能 | MTP、激活检查点、融合线性交叉熵 | | 内核 | DeepEP 分发 + torch_mm 专家 + TransformerEngine | | 指标 | NeMo AutoModel（EP=64） | | 每 GPU TPS（平均） | 815 | | 每 GPU TFLOP/s | ~293 | | 峰值内存 | 58.2 GiB |

为什么没有 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 + 自定义内核）。

8x H100 80GB（单节点）

关于路由门的说明。以下 NeMo AutoModel 的数字使用了平衡的路由门，它强制令牌在专家之间均匀分布。这模拟了 MoE 训练所追求的理想运行点：一个训练良好的模型的负载均衡损失会驱动专家利用率接近均匀，因此平衡路由反映了真实工作负载收敛到的稳态（并消除了随机虚拟令牌对专家并行性注入的“拖后腿”噪声）。v4/v5 在相同的虚拟令牌上运行其原生路由器。因此，平衡门衡量的是 NeMo AutoModel 在其目标 MoE 运行点的表现，而 v4/v5 的列则反映了它们的开箱即用行为。

Qwen3-30B-A3B

v4

v5 (FA2 + grouped_mm)

NeMo AutoModel (EP=8)

v5 → NeMo AutoModel

deadlock

3,075

11,340

3.69x

—

68.2 GiB

48.1 GiB

-29%

Avg Forward+Loss

582 ms

194 ms

3.00x

Avg Backward

758 ms

178 ms

4.26x

为什么 v4 会出现死锁：Transformers v4 将 Qwen3 MoE 专家存储为 128 个单独 MLP 模块的 ModuleList，每个模块都单独使用 FSDP 包裹。前向传播使用一个数据依赖的循环，仅迭代接收到令牌的专家。由于每个 rank 的数据不同，不同的 rank 会跳过不同的专家，导致 FSDP 的 AllGather/ReduceScatter 集体操作不匹配，从而引发无限期挂起。Transformers v5 通过将专家存储为融合的 3D 参数张量来修复此问题（没有每个专家的模块，也没有每个专家的 FSDP 集体操作）。

Nemotron 3 Nano 30B A3B

v4 (hub code)

v5 (FA2 + grouped_mm + Mamba CUDA)

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 安全的（无论令牌分配如何，都迭代所有专家），因此它不会像 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%），为更大的批量大小或更长的序列腾出了内存空间。

• DeepEP 将通信与计算融合。与专家路由的独立 AllGather/ReduceScatter 集体操作不同，DeepEP 融合了令牌分发，并将其合并为优化的 GPU 内核，从而在专家计算过程中重叠通信。

• TransformerEngine 内核加速了核心操作。TE 的融合注意力、线性层和 RMSNorm 实现提供了在所有层类型（不仅仅是 MoE 层）上一致的速度提升，相较于 PyTorch/Flash Attention 的等效实现。

HuggingFace AutoModel 所利用的 Transformers v5 特性

专家后端

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

后端

描述

最适合

eager

对选定专家进行循环

调试、兼容性和正确性。v4 也支持。

batched_mm

复制专家参数，通过 torch.bmm 进行单个批量 GEMM

小输入，使用 torch.compile 快速。v5 新增。

grouped_mm

按专家排序令牌，通过 torch.nn.functional.grouped_mm 进行单个分组 GEMM

训练（内存高效，无参数复制）。v5 新增。

grouped_mm 后端是关键的训练优化：它不再逐个循环专家，而是根据分配的专家对令牌进行排序，并执行一次融合的分组矩阵乘法。

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"
,
# 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 则路由令牌并使用 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 的专家。

# 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，它将令牌路由融合到优化的 GPU 内核中，并在与分组 GEMM 结合进行分组专家计算时，显著提高了速度。在我们的大规模 MoE 基准测试中，DeepEP + 分组 GEMM 相比于 all-gather + 循环专家基线，在完整的 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 倍的训练吞吐量，同时减少了 29-32% 的 GPU 内存使用。由于真正的专家并行性在 GPU 之间分片专家，相同的路径可以扩展到在 16 个节点上对 Nemotron 3 Ultra 这样的 550B 模型进行完整微调，这是专家并行性成为内存中容纳模型的关键的领域。由于 NeMo AutoModel 检查点是标准的 HF 格式 safetensors，因此可以在 vLLM 和 SGLang 等推理框架上部署它们。

代码、配置和基准脚本都可以在 NeMo AutoModel 仓库中找到。

致谢

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

本文中提到的模型 2

更多来自该作者的文章

Nemotron 3.5 内容安全：可定制的多模态安全，适用于全球企业 AI

12

2026 年 6 月 4 日

如何为您的语言、领域或口音微调 Nemotron 3.5 ASR

66

社区

编辑

预览

通过拖动到文本输入框、粘贴或

点击此处

上传图片、音频和视频。

轻点或粘贴此处以上传图片

评论

· 注册或登录以评论

• +14