返回
原文
Hugging Face：Blog（RSS）
精选
66AI 编辑部评分，满分 100

NVIDIA NeMo AutoModel：一行代码加速Transformer MoE模型微调

2026-06-25 00:00· 45天前
AI 导读

NVIDIA NeMo AutoModel 是基于 Transformers v5 的开源库，添加 Expert Parallelism、DeepEP 融合 all-to-all 调度和 TransformerEngine 内核。在 MoE 模型微调中，相比原生 v5，训练吞吐量提升 3.4–3.7 倍，GPU 内存减少 29–32%，仅需改动一行 import。在 16 节点 128 张 H100 上全微调 Nemotron 3 Ultra 550B A55B 时，v5 因内存不足无法运行，而 AutoModel 凭借 EP=64 专家并行使训练可行。单节点 30B MoE 模型（如 Qwen3-30B-A3B）同样获得可量化的性能优势。

推荐理由

英伟达的 NeMo AutoModel 把 MoE 模型微调速度提高了三倍多，内存省了近三分之一，代码只需改一行 import，做训练的可以立刻升级。

正文 · AI 翻译

v5 版本奠定了混合专家模型（MoE）的基础：专家后端、动态权重加载以及分布式执行，使 MoE 具备可扩展性且易于构建。

NVIDIA NeMo AutoModel 是一个开源库，属于 NVIDIA NeMo 框架的一部分，用于大规模构建自定义生成式 AI 模型。NeMo AutoModel 在 v5 之上进行了简洁的构建，增加了专家并行（Expert Parallelism）、DeepEP 融合的全对全（all-to-all）调度以及 TransformerEngine 内核，并利用 v5 的动态权重加载将这些优化应用于广泛且不断增长的模型族。其成果是，在使用相同的 `from_pretrained()` API 微调 MoE 模型时，训练吞吐量提升了 3.4-3.7 倍，GPU 内存减少了 29-32%：只需修改一行导入代码，无需其他任何代码更改。

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

背景

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

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

NeMo AutoModel 通过继承 AutoModelForCausalLM，并增加专家并行（EP）、DeepEP 融合的全对全调度以及 TransformerEngine 内核，在 v5 之上进行构建。DeepEP 是 v5 尚未具备的部分：它使通信与专家计算重叠。并且，由于 NeMo AutoModel 借助 v5 的可逆权重转换来加载每个模型，它可以将工程精力集中在这些可复用的核心操作上，而不是每个模型的检查点适配工作上，同时 `save_pretrained()` 仍然输出标准的 Hugging Face 检查点，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()

这通过一次 from_pretrained() 调用，即可实现 FSDP2、专家并行、TransformerEngine 内核和 DeepEP 调度带来的速度提升、可扩展性和内存优化。

性能对比

我们在两种场景下评估了 NeMo AutoModel：跨 16 个节点对前沿规模的 550B 模型进行全参数微调，以及在单个节点上训练两个 30B MoE 模型。550B 的结果说明了专家并行在大规模场景下的必要性；30B 的结果则量化了相比 Transformers v5 的每 GPU 速度提升。

Nemotron 3 Ultra 550B A55B（全参数微调，多节点）

Nemotron 3 Ultra 550B A55B 是一个 550B 参数的混合模型，集成了 Mamba2、LatentMoE 和多 token 预测（MTP）。我们对其进行了全参数微调基准测试：每个参数都参与更新，Adam 优化器状态被完整实例化，在此规模下需要跨越 16 个 H100 节点（128 个 GPU）。

方法：

参数 值
硬件 16 个 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 能够适配。

单节点 30B MoE 基准测试

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

方法：

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

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

Qwen3-30B-A3B

指标 v4 v5 (FA2 + grouped_mm) NeMo AutoModel (EP=8) v5 → NeMo AutoModel
每 GPU 每秒 token 数（平均） 死锁 3,075 11,340 3.69 倍
峰值内存 — 68.2 GiB 48.1 GiB -29%
平均前向+损失 — 582 毫秒 194 毫秒 3.00 倍
平均反向 — 758 毫秒 178 毫秒 4.26 倍

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（Hub 代码） v5（FA2 + grouped_mm + Mamba CUDA） NeMo AutoModel（EP=8） v5 → NeMo AutoModel
TPS/GPU（平均） 1,807 4,583 15,421 3.36 倍
峰值内存 61.9 GiB 62.1 GiB 42.5 GiB -32%
平均前向+损失 1,024 毫秒 283 毫秒 109 毫秒 2.60 倍
平均反向 1,246 毫秒 611 毫秒 157 毫秒 3.89 倍

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

加速的来源

NeMo AutoModel 相比 Transformers v5 的 3.4-3.7 倍加速来自三个方面：

  1. 专家并行降低了内存压力。EP=8 将专家权重分布到多个 GPU 上，将每个 GPU 的 MoE 内存占用降低 8 倍。对于 Qwen3，这使峰值内存从 68.2 GiB 降至 48.1 GiB（-29%）。对于 Nemotron Nano，则从 62.1 GiB 降至 42.5 GiB（-32%），为更大的批次大小或更长的序列腾出了空间。

  2. DeepEP 将通信与计算融合。DeepEP 没有为专家路由使用单独的 AllGather/ReduceScatter 集合通信，而是将 token 分发与合并融合到优化的 GPU 内核中，使通信与专家计算重叠进行。

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

HuggingFace AutoModel 利用的 Transformers v5 特性

专家后端

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-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",           
    linear="te",         
    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 上。


from torch.distributed.tensor import Shard, distribute_tensor


distribute_tensor(param, device_mesh, [Shard(0)])

在 8 块 GPU 上设置 ep_size=8 时，每块 GPU 仅持有 1/8 的专家参数。对于像 Nemotron-3-Nano-30B-A3B 这样专家权重约 55 GiB 的模型，专家并行将每块 GPU 的专家占用从约 55 GiB 降至约 6.8 GiB，从而使得仅靠 FSDP 方式会内存溢出的训练成为可能。

在专家并行基础上，NeMo AutoModel 集成了 DeepEP，它将 token 路由融合到优化的 GPU 内核中，并在结合分组 GEMM 进行分组专家计算时带来显著的加速效果。在我们的大规模 MoE 基准测试中，与 all-gather 加循环专家基线相比，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 性能总结
  • HuggingFace 上的 NeMo AutoModel

结论

NVIDIA NeMo AutoModel 是 HuggingFace 用户扩展模型训练规模的必然下一步。通过直接构建在 Transformers v5 之上，AutoModel 提供了一条零摩擦的升级路径：只需更改一行导入代码，即可获得一个速度提升超过三倍的模型实例。

在 Qwen3-30B-A3B 和 Nemotron 3 Nano 30B-A3B 上，与最佳 Transformers v5 配置相比，该方案实现了 3.4 至 3.7 倍的训练吞吐量提升，同时 GPU 内存占用减少 29% 至 32%。由于真正的专家并行机制将专家模型分布在多个 GPU 上，同样的路径可扩展至对 Nemotron 3 Ultra 这类 550B 模型进行全量微调，仅需 16 个节点——这正是专家并行对于将模型装入内存至关重要的场景。由于 NeMo AutoModel 检查点采用标准 HF 格式的 safetensors，你可以将其部署在 vLLM 和 SGLang 等推理框架上。

代码、配置文件和基准测试脚本均已发布在 NeMo AutoModel 仓库中。

致谢

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

开源/仓库教程/实践数据/训练

来源：Hugging Face：Blog（RSS） · huggingface.co

导出 Markdown
返回
精选66打开原文
导出 Markdown

NVIDIA NeMo AutoModel：一行代码加速Transformer MoE模型微调

Hugging Face：Blog（RSS）·2026-06-25 00:00·45天前
AI 导读

NVIDIA NeMo AutoModel 是基于 Transformers v5 的开源库，添加 Expert Parallelism、DeepEP 融合 all-to-all 调度和 TransformerEngine 内核。在 MoE 模型微调中，相比原生 v5，训练吞吐量提升 3.4–3.7 倍，GPU 内存减少 29–32%，仅需改动一行 import。在 16 节点 128 张 H100 上全微调 Nemotron 3 Ultra 550B A55B 时，v5 因内存不足无法运行，而 AutoModel 凭借 EP=64 专家并行使训练可行。单节点 30B MoE 模型（如 Qwen3-30B-A3B）同样获得可量化的性能优势。

正文 · AI 翻译

v5 版本奠定了混合专家模型（MoE）的基础：专家后端、动态权重加载以及分布式执行，使 MoE 具备可扩展性且易于构建。

NVIDIA NeMo AutoModel 是一个开源库，属于 NVIDIA NeMo 框架的一部分，用于大规模构建自定义生成式 AI 模型。NeMo AutoModel 在 v5 之上进行了简洁的构建，增加了专家并行（Expert Parallelism）、DeepEP 融合的全对全（all-to-all）调度以及 TransformerEngine 内核，并利用 v5 的动态权重加载将这些优化应用于广泛且不断增长的模型族。其成果是，在使用相同的 `from_pretrained()` API 微调 MoE 模型时，训练吞吐量提升了 3.4-3.7 倍，GPU 内存减少了 29-32%：只需修改一行导入代码，无需其他任何代码更改。

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

背景

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

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

NeMo AutoModel 通过继承 AutoModelForCausalLM，并增加专家并行（EP）、DeepEP 融合的全对全调度以及 TransformerEngine 内核，在 v5 之上进行构建。DeepEP 是 v5 尚未具备的部分：它使通信与专家计算重叠。并且，由于 NeMo AutoModel 借助 v5 的可逆权重转换来加载每个模型，它可以将工程精力集中在这些可复用的核心操作上，而不是每个模型的检查点适配工作上，同时 `save_pretrained()` 仍然输出标准的 Hugging Face 检查点，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()

这通过一次 from_pretrained() 调用，即可实现 FSDP2、专家并行、TransformerEngine 内核和 DeepEP 调度带来的速度提升、可扩展性和内存优化。

性能对比

我们在两种场景下评估了 NeMo AutoModel：跨 16 个节点对前沿规模的 550B 模型进行全参数微调，以及在单个节点上训练两个 30B MoE 模型。550B 的结果说明了专家并行在大规模场景下的必要性；30B 的结果则量化了相比 Transformers v5 的每 GPU 速度提升。

Nemotron 3 Ultra 550B A55B（全参数微调，多节点）

Nemotron 3 Ultra 550B A55B 是一个 550B 参数的混合模型，集成了 Mamba2、LatentMoE 和多 token 预测（MTP）。我们对其进行了全参数微调基准测试：每个参数都参与更新，Adam 优化器状态被完整实例化，在此规模下需要跨越 16 个 H100 节点（128 个 GPU）。

方法：

参数 值
硬件 16 个 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 能够适配。

单节点 30B MoE 基准测试

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

方法：

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

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

Qwen3-30B-A3B

指标 v4 v5 (FA2 + grouped_mm) NeMo AutoModel (EP=8) v5 → NeMo AutoModel
每 GPU 每秒 token 数（平均） 死锁 3,075 11,340 3.69 倍
峰值内存 — 68.2 GiB 48.1 GiB -29%
平均前向+损失 — 582 毫秒 194 毫秒 3.00 倍
平均反向 — 758 毫秒 178 毫秒 4.26 倍

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（Hub 代码） v5（FA2 + grouped_mm + Mamba CUDA） NeMo AutoModel（EP=8） v5 → NeMo AutoModel
TPS/GPU（平均） 1,807 4,583 15,421 3.36 倍
峰值内存 61.9 GiB 62.1 GiB 42.5 GiB -32%
平均前向+损失 1,024 毫秒 283 毫秒 109 毫秒 2.60 倍
平均反向 1,246 毫秒 611 毫秒 157 毫秒 3.89 倍

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

加速的来源

NeMo AutoModel 相比 Transformers v5 的 3.4-3.7 倍加速来自三个方面：

  1. 专家并行降低了内存压力。EP=8 将专家权重分布到多个 GPU 上，将每个 GPU 的 MoE 内存占用降低 8 倍。对于 Qwen3，这使峰值内存从 68.2 GiB 降至 48.1 GiB（-29%）。对于 Nemotron Nano，则从 62.1 GiB 降至 42.5 GiB（-32%），为更大的批次大小或更长的序列腾出了空间。

  2. DeepEP 将通信与计算融合。DeepEP 没有为专家路由使用单独的 AllGather/ReduceScatter 集合通信，而是将 token 分发与合并融合到优化的 GPU 内核中，使通信与专家计算重叠进行。

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

HuggingFace AutoModel 利用的 Transformers v5 特性

专家后端

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-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",           
    linear="te",         
    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 上。


from torch.distributed.tensor import Shard, distribute_tensor


distribute_tensor(param, device_mesh, [Shard(0)])

在 8 块 GPU 上设置 ep_size=8 时，每块 GPU 仅持有 1/8 的专家参数。对于像 Nemotron-3-Nano-30B-A3B 这样专家权重约 55 GiB 的模型，专家并行将每块 GPU 的专家占用从约 55 GiB 降至约 6.8 GiB，从而使得仅靠 FSDP 方式会内存溢出的训练成为可能。

在专家并行基础上，NeMo AutoModel 集成了 DeepEP，它将 token 路由融合到优化的 GPU 内核中，并在结合分组 GEMM 进行分组专家计算时带来显著的加速效果。在我们的大规模 MoE 基准测试中，与 all-gather 加循环专家基线相比，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 性能总结
  • HuggingFace 上的 NeMo AutoModel

结论

NVIDIA NeMo AutoModel 是 HuggingFace 用户扩展模型训练规模的必然下一步。通过直接构建在 Transformers v5 之上，AutoModel 提供了一条零摩擦的升级路径：只需更改一行导入代码，即可获得一个速度提升超过三倍的模型实例。

在 Qwen3-30B-A3B 和 Nemotron 3 Nano 30B-A3B 上，与最佳 Transformers v5 配置相比，该方案实现了 3.4 至 3.7 倍的训练吞吐量提升，同时 GPU 内存占用减少 29% 至 32%。由于真正的专家并行机制将专家模型分布在多个 GPU 上，同样的路径可扩展至对 Nemotron 3 Ultra 这类 550B 模型进行全量微调，仅需 16 个节点——这正是专家并行对于将模型装入内存至关重要的场景。由于 NeMo AutoModel 检查点采用标准 HF 格式的 safetensors，你可以将其部署在 vLLM 和 SGLang 等推理框架上。

代码、配置文件和基准测试脚本均已发布在 NeMo AutoModel 仓库中。

致谢

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

来源：Hugging Face：Blog（RSS）· huggingface.co