NVIDIA NeMo 加速 Transformer 模型微调训练

NVIDIA NeMo AutoModel 是 NVIDIA NeMo 框架的开源库，用于大规模构建定制生成式 AI 模型。NeMo AutoModel 在 Transformers v5 的基础上进行了扩展，增加了 Expert Parallelism（专家并行）、DeepEP 融合全对全调度和 TransformerEngine 内核，并利用 v5 的动态权重加载为更广泛的模型系列提供这些优化。效果是在微调 MoE 模型时，使用相同的 from_pretrained() API，训练吞吐量提高 3.4-3.7 倍，GPU 内存减少 29-32%，相比原生 Transformers v5，只需改一行导入语句，无需其他代码修改。

本文详细介绍了这个方案的工作原理，以及用户如何在不改变 API 的情况下更快地微调 MoE 模型。

MoE 模型训练的挑战

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

Transformers v5（"v5"）引入了一级 MoE 支持，包括专家后端、动态权重加载和用于分布式执行的张量并行计划。此外，v5 通过直接将 PyTorch 的 DeviceMesh 集成到 from_pretrained() 中，使分布式训练成为一级特性。

NeMo AutoModel 通过继承 AutoModelForCausalLM 并加入 Expert Parallelism（EP）、DeepEP 融合全对全调度和 TransformerEngine 内核，在 v5 的基础上进行构建。DeepEP 是 v5 还没有的：它重叠通信和专家计算。由于 NeMo AutoModel 使用 v5 的可逆权重转换来加载每个模型，它可以把工程重点放在这些可复用的核心操作上，而不是逐模型的检查点管道，同时 save_pretrained() 仍然输出标准 HF 检查点，vLLM 和 SGLang 等工具可以加载。

下一部分讲解两者如何协同工作及我们测得的性能收益，从全微调 NVIDIA Nemotron 3 Ultra 550B A55B 跨 16 个节点，一直到 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 训练。要用 Expert Parallelism 在 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、Expert Parallelism、TransformerEngine 内核和 DeepEP 调度，全部来自 from_pretrained() 调用。

性能对比

我们在两个体制下评估了 NeMo AutoModel：全微调跨 16 个节点的前沿规模 550B 模型，以及在单个节点上训练两个 30B MoE 模型。550B 的结果展示了为什么 Expert Parallelism 在大规模训练中至关重要；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 基准测试

我们在配有 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 列反映了其开箱即用的行为。

为什么 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 倍加速来自三个来源：

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%），释放空间用于更大的批次大小或更长序列。

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%），释放空间用于更大的批次大小或更长序列。

DeepEP 融合通信与计算。 DeepEP 不是为专家路由分别执行 AllGather/ReduceScatter 集合，而是融合 token 调度并合并成优化的 GPU 内核，重叠通信与专家计算。

DeepEP 融合通信与计算。 DeepEP 不是为专家路由分别执行 AllGather/ReduceScatter 集合，而是融合 token 调度并合并成优化的 GPU 内核，重叠通信与专家计算。

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

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

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
)

Expert Parallelism 和 DeepEP

Transformers v5 也附带了一条 Expert Parallelism 路径。它在 GPU 间分片专家权重。GroupedGemmParallel 风格只加载每个设备的本地专家，RouterParallel 路由 token 并用全规约组合结果。它整洁地构建在 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 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 基准测试中，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 性能总结

NeMo AutoModel 在 HuggingFace 上

对于扩展模型训练的 HuggingFace 用户，NVIDIA NeMo AutoModel 是自然的下一步。通过直接构建在 Transformers v5 之上，AutoModel 提供了零摩擦的升级路径：改一行导入并获得一个快三倍多的模型实例。

在 Qwen3-30B-A3B 和 Nemotron 3 Nano 30B-A3B 上，相比最佳的 Transformers v5 配置，这提供了 3.4-3.7 倍更高的训练吞吐量，GPU 内存减少 29-32%。并且由于真正的 Expert Parallelism 在 GPU 间分片专家，相同的路径扩展到全微调像 Nemotron 3 Ultra 这样的 550B 模型跨 16 个节点，这是 Expert Parallelism 成为将模型放入内存的必需条件的体制。因为 NeMo AutoModel 检查点是标准 HF 格式的 safetensors，你可以在 vLLM 和 SGLang 这样的推理框架上部署它们。

代码、配置和基准脚本都可在 NeMo AutoModel 仓库中获得。

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