加速 Transformer 微调：NVIDIA NeMo AutoModel 深度解读

NVIDIA NeMo AutoModel 是 NVIDIA NeMo 框架的一部分，这是一个用于大规模构建自定义生成式 AI 模型的开源库。NeMo AutoModel 建立在 Transformers v5 之上，引入了专家并行（Expert Parallelism）、DeepEP 融合的全对全（all-to-all）分发以及 TransformerEngine 内核等优化。它利用 v5 的动态权重加载功能，将这些优化扩展到一个广泛且不断增长的模型家族中。

其结果是，与原生 Transformers v5 相比，在微调混合专家（MoE）模型时，训练吞吐量提高了 3.4-3.7 倍，GPU 内存减少了 29-32%，且无需更改代码，仅通过一行导入语句即可使用相同的 from_pretrained() API。

背景

MoE（Mixture of Experts，混合专家）模型的兴起为高效训练带来了新的挑战：将令牌（tokens）路由到数百个专家、将专家矩阵乘法融合为单个内核、在 GPU 间分片权重，以及将通信与计算重叠，这些都需要超出通用库默认提供的基础设施支持。

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

NeMo AutoModel 通过子类化 AutoModelForCausalLM 并在 v5 的基础上添加专家并行（EP）、DeepEP 融合的全对全分发以及 TransformerEngine 内核，进一步增强了 v5 的能力。DeepEP 是 v5 尚未具备的关键组件，它实现了通信与专家计算的 overlap（重叠）。由于 NeMo AutoModel 依托 v5 的可逆权重转换来加载每个模型，它可以将工程重点放在这些可复用的核心操作上，而不是针对每个模型的检查点管道。同时，save_pretrained() 依然输出标准的 Hugging Face 检查点，使得 vLLM 和 SGLang 等工具能够加载这些模型。

核心内容

API 兼容性与无缝集成

NeMo AutoModel 的主要目标之一是保持与 Hugging Face Transformers 的 API 兼容性，以赋能开源社区。NeMoAutoModelForCausalLM 子类化自 AutoModelForCausalLM，因此任何适用于 HF 模型的代码同样适用于 AutoModel。

用户只需更改一行导入代码即可切换：

# 原生 Hugging Face
from transformers import AutoModelForCausalLM

# NVIDIA NeMo AutoModel
from nemo_automodel import NeMoAutoModelForCausalLM

这一行导入背后做了大量工作：

1. 针对流行架构的手动调优：对于 Qwen3、NVIDIA Nemotron、GPT-OSS 和 DeepSeek V3 等流行的 MoE 架构，NeMo AutoModel 提供了经过手工调优的实现，包括 TransformerEngine 注意力机制、融合线性层和自定义专家内核。
2. 通用回退机制：对于其他模型，它回退到原生 HF，但仍应用诸如 Liger 内核修补等优化措施。
3. 多 GPU 扩展性：无论走哪条路径，生成的模型都准备好进行扩展：只需传入 device_mesh，即可实现多 GPU 训练，无需进一步重写代码。

多 GPU 训练的实现示例

NeMo AutoModel 的真正优势在于将 MoE 模型扩展到多 GPU 训练。以下是在 8 张 GPU 上使用专家并行（EP=8）训练 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 分发，从而带来速度、可扩展性和内存优化。

性能基准测试

评估分为两个场景：在 16 个节点上全微调一个前沿规模的 550B 模型，以及在单个节点上训练两个 30B MoE 模型。

1. Nemotron 3 Ultra 550B A55B（全微调，多节点）

Nemotron 3 Ultra 550B A55B 是一个拥有 5500 亿参数的混合模型，包含 Mamba2、LatentMoE 和多令牌预测（MTP）。基准测试涉及全微调，即更新每个参数并实例化 Adam 优化器状态。这在 16 个 H100 节点（共 128 张 GPU）上进行。

• 为何没有 Transformers v5 的数据？ 在如此大的规模下，Transformers v5 会因内存不足（OOM）而崩溃，因此无法提供 v5 的数据。
• 专家并行的作用：AutoModel 的专家并行将专家权重分片到不同的 GPU 上，将内存占用控制在预算范围内，从而使得全微调成为可能。

2. 单节点 30B MoE 基准测试

在配备 8 张 H100 80GB GPU 的单节点上，基准测试了三种方法：

1. HF Transformers v4（Hub 代码）
2. HF Transformers v5（使用最佳可用优化）
3. NeMo AutoModel（EP=8 + 自定义内核）

关于路由门控（Routing Gate）的说明： NeMo AutoModel 的基准数据使用平衡路由门控，强制令牌均匀分布在专家之间。这模拟了 MoE 训练趋向的理想操作点：训练良好的模型其负载均衡损失会将专家利用率驱动至接近均匀，因此平衡路由反映了真实工作负载收敛到的稳态，并消除了随机虚拟令牌引入专家并行中的“落后者”噪声。v4/v5 则在相同的虚拟令牌上运行其原生路由器。因此，平衡门控衡量的是 NeMo AutoModel 在其目标 MoE 操作点上的表现，而 v4/v5 列反映的是其开箱即用的行为。

具体案例分析：

• Qwen3-30B-A3B：

  • v4 死锁原因：Transformers v4 将 Qwen3 MoE 专家存储为 ModuleList，其中包含 128 个独立的 MLP 模块，每个模块单独包裹 FSDP。前向传播使用数据依赖循环，仅迭代接收到令牌的专家。由于不同 Rank 的数据不同，不同 Rank 跳过的专家也不同，导致不匹配的 FSDP AllGather/ReduceScatter 集合操作，造成无限期挂起。
  • v5 的修复：v5 通过将专家存储为融合的 3D 参数张量（无每个专家的独立模块，无每个专家的 FSDP 集合操作）解决了此问题。

• Nemotron 3 Nano 30B A3B：

  • v4 配置：使用 trust_remote_code=True（NVIDIA 的 Hub 建模代码）。由于 Hub 代码中的专家循环是 FSDP 安全的（无论令牌分配如何都迭代所有专家），因此不会像 Qwen3 v4 那样死锁。

性能提升的来源

NeMo AutoModel 相比 Transformers v5 的 3.4-3