Accelerating 成为 AI 产业观察信号

NVIDIA NeMo AutoModel 是 NVIDIA NeMo 框架中的一个开源库，用于大规模构建定制生成式 AI 模型。NeMo AutoModel 在 v5 之上进行了干净的封装，新增了 Expert Parallelism、DeepEP fused all-to-all dispatch 和 TransformerEngine kernels，并借助 v5 的动态权重加载，将这些优化带给更广泛且仍在扩展的模型家族。其结果是：在 MoE 模型微调中，训练吞吐量比原生 Transformers v5 提升 3.4-3.7 倍，GPU 内存占用减少 29-32%，且仍使用同一个 from_pretrained() API——只需一行 import，代码无需做其他改动。

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

MoE 模型的兴起给高效训练带来了新的挑战：在数百个专家之间路由 token、把专家矩阵乘法融合进单个 kernel、在 GPU 之间切分权重，以及让通信与计算重叠，这些都需要超出通用库开箱即用能力的基础设施支持。

Transformers v5（“v5”）引入了原生的 MoE 支持，例如 expert backends、dynamic weight loading，以及用于分布式执行的 tensor parallel plans。此外，v5 还通过将 PyTorch 的 DeviceMesh 直接集成到 from_pretrained() 中，让分布式训练成为一等能力。

NeMo AutoModel 在 v5 的基础上构建，通过继承 `AutoModelForCausalLM`，并加入专家并行（EP）、DeepEP 融合式 all-to-all 派发以及 TransformerEngine 内核。DeepEP 是 v5 目前还不具备的部分：它将通信与专家计算重叠执行。由于 NeMo AutoModel 依托 v5 的可逆权重转换来加载每个模型，它可以把工程重点放在这些可复用的核心操作上，而不必为每个模型单独处理 checkpoint 管线；与此同时，`save_pretrained()` 仍然会导出标准的 HF checkpoints，像 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 attention、融合线性层和自定义专家 kernel。对于其他模型，它会回退到原生 HF，同时仍会应用诸如 Liger kernel patching 等优化。无论走哪条路径，生成的模型都已准备好扩展：传入一个 device_mesh，就可以在无需进一步重写的情况下进行多 GPU 训练。

NeMo AutoModel 真正出色的地方，在于将 MoE 模型扩展到多 GPU 训练。要在 8 块 GPU 上使用 Expert Parallelism 训练 Nemotron 3 Nano 30B A3B，只需添加分布式 mesh 配置：

这带来了速度、可扩展性和内存优化，借助 FSDP2、Expert Parallelism、TransformerEngine kernels 和 DeepEP dispatch，全部都来自一次 `from_pretrained()` 调用。

我们在两种场景下评估了 NeMo AutoModel：一种是在 16 个节点上对一个前沿规模的 550B 模型进行全量微调；另一种是在单节点上训练两个 30B MoE 模型。550B 的结果说明了 Expert Parallelism 在大规模场景下为何至关重要；30B 的结果则量化了相较 Transformers v5 的每 GPU 加速收益。

Nemotron 3 Ultra 550B A55B 是一个拥有 550B 参数的混合模型，随 Mamba2、LatentMoE 和 Multi-Token Prediction（MTP）一起发布。我们对其进行了完整微调基准测试：每个参数都会被更新，Adam 优化器状态也会被完整实例化；在这个规模下，这一过程横跨 16 个 H100 节点（128 张 GPU）。

为什么没有 Transformers v5 这一列。Transformers v5 在这个规模下会发生显存耗尽，因此这里没有可报告的 v5 数值。AutoModel 的 Expert Parallelism 会把专家分片到各个 GPU 上，将占用控制在预算范围内，这正是完整微调能够运行的原因。下面的 30B 对比展示了在 v5 能够适用时，同样的优势。

我们在单节点、8x H100 80GB GPU 上对三种方案进行了基准测试：HF Transformers v4（hub 代码）、HF Transformers v5（采用可用的最佳优化），以及 NeMo AutoModel（EP=8 + 自定义 kernel）。

关于路由 gate 说明一下。下面的 NeMo AutoModel 数值使用的是平衡路由 gate，它会强制将 token 均匀分配到各个 expert 上。这模拟了 MoE 训练所追求的理想运行点：训练良好的模型，其负载均衡损失会把 expert 利用率驱动到接近均匀，因此平衡路由反映的是实际工作负载最终收敛到的稳态（同时也移除了随机 dummy token 通过 expert parallelism 注入的 straggler 噪声）。v4/v5 在相同的 dummy token 上运行其原生 router。因此，平衡 gate 测量的是 NeMo AutoModel 在其目标 MoE 运行点下的表现，而 v4/v5 列则反映它们开箱即用的行为。

v4 死锁的原因：Transformers v4 将 Qwen3 的 MoE 专家存储为一个包含 128 个独立 MLP 模块的 `ModuleList`，每个模块都分别用 FSDP 包装。前向过程使用了一个依赖数据的循环，只会遍历接收到 token 的专家。由于不同 rank 上的数据不同，不同的 rank 会跳过不同的专家，从而导致 FSDP 的 `AllGather`/`ReduceScatter` 集合通信不匹配，进而无限挂起。Transformers v5 通过将专家存储为融合的 3D 参数张量来修复这一问题（不再有按专家拆分的模块，也就不再有按专家拆分的 FSDP 集合通信）。

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 上，使单卡上的 MoE 占用缩小 8 倍。对于 Qwen3，这会把峰值内存从 68.2 GiB 降到 48.1 GiB（-29%）。对于 Nemotron Nano，则从 62.1 GiB 降到 42.5 GiB（-32%），为更大的 batch size 或更长的序列留出余量。

DeepEP 将通信与计算融合在一起。它不再为专家路由分别执行 AllGather/ReduceScatter 这类集合通信，而是把 token 派发和合并融合到经过优化的 GPU 内核中，让通信与专家计算重叠进行。

TransformerEngine 内核加速核心操作。TE 的融合 attention、线性层和 RMSNorm 实现，相比 PyTorch/Flash Attention 对应实现，在所有层类型上都能带来稳定提速，而不仅仅是 MoE 层。

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

`grouped_mm` 后端是关键的训练优化：它不再逐个 expert 依次循环，而是先按 token 所分配的 expert 对其排序，再执行一次融合的 grouped matrix multiplication。

NeMo AutoModel 进一步推进了这一点。对于带有自定义实现的模型，它会使用 DeepEP 融合式 all-to-all 分发，结合 grouped GEMM 内核和 TransformerEngine 线性层。其演进过程如下：

在 NeMo AutoModel 中，expert backend 通过 BackendConfig 进行配置：

Transformers v5 也提供了一条 Expert Parallelism 路径。它会将 expert 权重切分到各个 GPU 上。GroupedGemmParallel 这种方式只加载每个设备本地的 experts，而 RouterParallel 则负责路由 token，并通过 all_reduce 汇总结果。它建立在 v5 现有的 tensor-parallel 机制之上，结构相当清晰。启用后，模型的 tp_plan 会返回其 expert plan，因此 expert parallelism 会与 data parallelism 共享设备预算（ep × dp = world_size）。在这里的单机 30B 基准测试中，我们发现纯 data-parallel 的 v5（dp=8，ep=1）是最快的 v5 配置，所以报告的也是这一套 v5 设置。

NeMo AutoModel 则采取了互补的方案，专门针对多 GPU MoE 训练进行优化。它把 EP 作为独立的并行维度，即一个专门的 moe_mesh，而不是从 data-parallel mesh 中切分出来，并使用 PyTorch 的 DTensor 和 Shard(0)。由于 expert mesh 与 data parallelism 是正交的，两者可以在同一批设备上叠加使用。在 8 张 GPU 上，NeMo AutoModel 会同时运行 ep=8 和 dp=8，因此每张 GPU 一边训练自己的数据切片，一边只持有 1/8 的 experts。expert 权重会沿 expert 维度在 GPU 之间物理切分。

在 8 张 GPU 上设置 `ep_size=8` 时，每张 GPU 只持有 1/8 的专家参数。以 Nemotron-3-Nano-30B-A3B 这类模型为例，其专家权重约为 55 GiB，EP 将单卡专家占用从约 55 GiB 降到约 6.8 GiB，使得原本仅靠 FSDP 容易耗尽显存的训练成为可能。

在 EP 之上，NeMo AutoModel 集成了 DeepEP，将 token 路由融合进经过优化的 GPU kernel，并在与 grouped GEMM 结合进行分组专家计算时带来显著加速。在我们的大规模 MoE 基准测试中，与 all-gather + 循环式专家基线相比，DeepEP + grouped GEMM 将完整的 DeepSeek V3 671B 模型的单次迭代成本降低了 47%。

Transformers v5 还通过 WeightConverter 和 WeightRenaming 引入了动态权重加载系统。这使得 MoE checkpoint 可以存储为融合后的 3D tensor，从而实现更高效的执行。WeightConverter 会在 `from_pretrained()` 期间应用可组合操作，按需转换 checkpoint tensor。

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 性能总结

NVIDIA NeMo AutoModel 是 HuggingFace 用户在扩展模型训练规模时的自然下一步。基于 Transformers v5 直接构建，AutoModel 提供了一条几乎零摩擦的升级路径：只需改一行 import，就能得到一个速度快出三倍以上的模型实例。

在 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 的 checkpoint 采用标准 HF-format safetensors，你可以将其部署到 vLLM 和 SGLang 这类推理框架上。

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

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

· 登录或注册以评论