文本生成 • 32B • 已更新 • 1.06M • 780 https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16

NVIDIA NeMo AutoModel 是 NVIDIA NeMo 框架 的一部分，是一个用于大规模构建自定义生成式 AI 模型的开源库。NeMo AutoModel 在 v5 之上构建得十分干净，新增了专家并行（Expert Parallelism）、DeepEP 融合的全对全调度（fused all-to-all dispatch）以及 TransformerEngine 内核，并借助 v5 的动态权重加载（dynamic weight loading）将这些优化应用于日益广泛和增长的模型家族。其回报是：在对 MoE 模型进行微调时，相比原生 Transformers v5，训练吞吐量提升 3.4-3.7 倍，GPU 内存减少 29-32%，且使用相同的 from_pretrained() API：只需更改一行导入代码，无需其他代码修改。

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

MoE 模型的兴起为高效训练带来了新的挑战：跨数百个专家路由 token、将专家矩阵乘法融合为单个内核、在 GPU 间分片权重，以及将通信与计算重叠——所有这些都需要通用库开箱即用之外的基础设施。

Transformers v5（“v5”）引入了顶级的 MoE 支持，例如专家后端、动态权重加载以及用于分布式执行的张量并行（tensor parallel）方案。此外，v5 通过将 PyTorch 的 DeviceMesh 直接集成到 from_pretrained() 中，使分布式训练成为了一等公民。

NeMo AutoModel 通过继承 AutoModelForCausalLM，并增加专家并行（EP）、DeepEP 融合的全对全调度以及 TransformerEngine 内核，从而在 v5 之上构建。DeepEP 是 v5 尚不具备的部分：它将通信与专家计算重叠。由于 NeMo AutoModel 借助 v5 的可逆权重转换（reversible weight conversion）来加载每个模型，它可以将工程精力集中于这些可复用的核心操作上，而非逐个模型的检查点适配，而 save_pretrained() 仍会输出标准的 HuggingFace 检查点，vLLM 和 SGLang 等工具可直接加载。

下一节将阐述两者如何协同工作，以及我们测量到的性能提升——涵盖从在 16 个节点上全量微调 NVIDIA Nemotron 3 Ultra 550B A55B，到单节点模型（如 Qwen3-30B-A3B 和 Nemotron 3 Nano 30B A3B）。

NeMo AutoModel 的目标之一是与 HuggingFace Transformers 保持 API 兼容，以支持开源社区。NeMoAutoModelForCausalLM 继承自 AutoModelForCausalLM，因此任何适用于 HF 模型的代码也同样适用于 AutoModel。

以下是在两者中加载模型的方式。仅导入语句发生了变化：

仅仅一个导入语句就完成了大量工作。对于 Qwen3、NVIDIA Nemotron、GPT-OSS 和 DeepSeek V3 等流行的 MoE 架构，NeMo AutoModel 提供了手工调优的实现，包含 TransformerEngine 注意力、融合线性层和自定义专家内核。对于其他架构，它会回退到原生 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 是一个 550B 参数的混合模型，集成了 Mamba2、LatentMoE 和多 Token 预测（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 每秒 token 数（平均）	815
每 GPU TFLOP/秒	~293
峰值内存	58.2 GiB

为何没有 Transformers v5 列。 Transformers v5 在此规模下会耗尽内存，因此这里无法报告 v5 数据。AutoModel 的专家并行将专家分片到多个 GPU 上，从而将占用空间控制在预算范围内，使得全量微调得以运行。下面的 30B 对比同样展示了 v5 可运行时 AutoModel 的优势。

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

方法论：

参数	值
硬件	8x H100 80GB（单个节点）
序列长度	4,096
本地批次大小	1

关于路由门的说明。 下面的 NeMo AutoModel 数据使用了均衡路由门（balanced routing gate），它强制 token 在专家之间均匀分布。这模拟了 MoE 训练所追求的理想运行点：一个训练良好的模型的负载均衡损失会使专家利用率接近均匀，因此均衡路由反映了实际工作负载收敛到的稳态（并消除了随机虚拟 token 可能注入专家并行中的掉队者噪声）。v4/v5 在相同的虚拟 token 上运行其原生路由器。因此，均衡门衡量的是 NeMo AutoModel 在其目标 MoE 运行点上的表现，而 v4/v5 列反映了它们的开箱即用行为。

指标	v4	v5 (FA2 + grouped_mm)	NeMo AutoModel (EP=8)	v5 → NeMo AutoModel
每 GPU 每秒 token 数（平均）	死锁	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 集合）解决了这个问题。

指标	v4（中心代码）	v5（FA2 + grouped_mm + Mamba CUDA）	NeMo AutoModel (EP=8)	v5 → NeMo AutoModel
每 GPU 每秒 token 数（平均）	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%），为更大的批次大小或更长的序列腾出了空间。

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

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

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 融合的全对全调度，结合分组 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 注意力
linear="te",         # TransformerEngine 线性层
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 上。

# 来自 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 基准测试中，在全 DeepSeek V3 671B 模型上，与 all-gather + 循环专家基线相比，DeepEP + 分组 GEMM 将每次迭代的成本降低了 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，请访问我们的官方文档页面以开始使用。

更多详细信息，请参见：

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。

来自该作者的更多内容