NVIDIA NeMo AutoModel 加速 Mixture-of-Experts 模型的微调

背景

Mixture-of-Experts 模型的兴起带来了训练挑战，Transformers v5 通过专家后端、动态权重加载和分布式执行来应对这些挑战，但仍缺少融合通信内核。

NeMo AutoModel：相同 API，更多性能

NeMo AutoModel 继承自 AutoModelForCausalLM，因此用户只需更改一行 import 即可从 Hugging Face Transformers 切换到 NeMo AutoModel，并在不修改训练代码的情况下获得 Expert Parallelism、DeepEP 融合 all‑to‑all 调度以及 TransformerEngine 内核。

加载模型时，NeMo AutoModel 与 Hugging Face 版本几乎相同，只是 import 不同：

from nemo_automodel import NeMoAutoModelForCausalLM
model = NeMoAutoModelForCausalLM.from_pretrained(
    "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
    dtype=torch.bfloat16,
)

对于 Qwen3、NVIDIA Nemotron、GPT‑OSS、DeepSeek V3 等流行的 MoE 架构，NeMo AutoModel 提供了使用 TransformerEngine attention、融合线性层和自定义专家内核的手工调优实现。对于其他模型，它回退到原生 Hugging Face，同时仍会应用诸如 Liger 内核打补丁等优化。通过传入 device_mesh，得到的模型可以扩展到多 GPU 训练。

要在 8 GPU 上使用 Expert Parallelism 训练 Nemotron 3 Nano 30B A3B，用户需要添加分布式 mesh 配置：

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、Expert Parallelism、TransformerEngine 内核和 DeepEP 调度的速度、可扩展性和内存优化。

性能比较

基准测试显示，NeMo AutoModel 在每 GPU 的 tokens‑per‑second 上提升 3.4‑3.7 倍，峰值内存降低 29‑32%，相较于最佳的 Transformers v5 配置，无论是在 16 节点的 550B 模型，还是单节点的两个 30B MoE 模型上均如此。

Nemotron 3 Ultra 550B A55B（完整微调，多节点）

Expert Parallelism 使得在 16 台 H100 节点上对 550B 参数的 Nemotron 3 Ultra 模型进行完整微调成为可能，而在该规模下 Transformers v5 会因内存不足而失败。

方法论：

硬件：16× H100 80GB（128 GPU）
专家并行：EP=64
本地批大小：2
序列长度：4,096
特性：MTP、激活检查点、融合线性交叉熵
内核：DeepEP 调度 + torch_mm 专家 + TransformerEngine

结果：

指标	NeMo AutoModel (EP=64)
TPS/GPU（平均）	815
TFLOP/s/GPU	~293
峰值内存	58.2 GiB

Transformers v5 在此规模下会因内存不足而无法运行，故未报告 v5 数据。

单节点 30B MoE 基准测试

在单个 8‑GPU H100 节点上，NeMo AutoModel 在吞吐量上比 Transformers v5 快 3.36‑3.69 倍，并将 Qwen3‑30B‑A3B 与 Nemotron 3 Nano 30B‑A3B 的峰值内存降低 29‑32%。

方法论：

硬件：8× H100 80GB（单节点）
序列长度：4,096
本地批大小：1

Qwen3-30B-A3B

指标	v4	v5 (FA2 + grouped_mm)	NeMo AutoModel (EP=8)	v5 → NeMo AutoModel
TPS/GPU（平均）	deadlock	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

Transformers v4 因将 Qwen3 MoE 专家存为 128 个单独的 MLP 模块的 ModuleList，每个模块单独进行 FSDP 包装，导致集合操作不匹配而死锁。Transformers v5 通过将专家存为融合的 3D 参数张量解决了此问题。

Nemotron 3 Nano 30B A3B

指标	v4 (hub code)	v5 (FA2 + grouped_mm + Mamba CUDA)	NeMo AutoModel (EP=8)	v5 → NeMo AutoModel
TPS/GPU（平均）	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 建模代码），其专家循环会遍历所有专家而不考虑 token 分配，从而避免了 Qwen3 v4 中出现的死锁。

加速来源

3.4‑3.7 倍的加速来源于 Expert Parallelism 对专家权重的分片、DeepEP 将 token 路由与计算融合，以及 TransformerEngine 内核对 attention 与线性层的加速。

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 融合通信与计算：不再使用独立的 AllGather/ReduceScatter 进行专家路由，DeepEP 将 token 调度与计算合并为优化的 GPU 内核，实现通信与计算的重叠。
TransformerEngine 内核加速核心算子：TE 的融合 attention、线性层和 RMSNorm 实现相较于原生 PyTorch/Flash Attention 在所有层（不仅限于 MoE 层）均提供一致的加速。

Transformers v5 功能在 HuggingFace AutoModel 中的利用

NeMo AutoModel 基于 Transformers v5 的专家后端（尤其是 grouped_mm）、动态权重加载和张量并行的 Expert Parallelism，并在其上额外加入 DeepEP 与 TransformerEngine。

专家后端

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

后端	描述	最佳适用
eager	对选定专家进行 for‑loop 迭代	调试、兼容性和正确性。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 融合 all‑to‑all 调度结合分组 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
)

专家并行和 DeepEP

Transformers v5 也提供了将专家权重在 GPU 之间分片的 Expert Parallelism 路径。NeMo AutoModel 采用互补的方式，专为多 GPU MoE 训练调优。它将 EP 设为独立的并行维度，构建专用的 moe_mesh（而非从 data‑parallel mesh 中切分），使用 PyTorch 的 DTensor 与 Shard(0)。由于专家 mesh 与数据并行是正交的，两者可以在同一设备上共同工作。在 8 GPU 上，NeMo AutoModel 同时运行 ep=8 与 dp=8，每个 GPU 只持有 1/8 的专家权重。

# 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)])

当 ep_size=8 且使用 8 GPU 时，每个 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 进行分组专家计算时提供显著的加速。

动态权重加载

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 格式的检查点，任何下游工具均可加载。

入门指南

用户可以按照 NVIDIA 官方文档安装 NeMo AutoModel，并通过一次 import 更改使用相同的 from_pretrained() API。

更多细节请参见：

NeMo AutoModel HuggingFace API Compatibility Guide
NeMo AutoModel Model Coverage
NeMo AutoModel Performance Summary
NeMo AutoModel on HuggingFace

结论

NeMo AutoModel 为希望在 MoE 训练上扩展的 Hugging Face 用户提供了零摩擦的升级路径，带来多倍的加速和显著的内存节省，同时生成兼容 vLLM 与 SGLang 的标准 Hugging Face 检查点。

代码、配置以及基准脚本均已在 NeMo AutoModel 仓库中公开。

SUMMARY: NVIDIA NeMo AutoModel 提供 3.4‑3.7 倍更高的训练吞吐量和 29‑32% 更低的 GPU 内存占用，用于微调 Mixture-of-Experts 模型。它基于 Hugging Face Transformers v5，加入 Expert Parallelism、DeepEP 与 TransformerEngine 内核，仅需一次 import 更改即可使用。

TITLE: NVIDIA NeMo AutoModel 加速 Mixture-of-Experts 模型的微调

NVIDIA NeMo AutoModel 加速 Mixture-of-Experts 模型的微调

NVIDIA NeMo AutoModel 加速 Mixture-of-Experts 模型的微调

背景

NeMo AutoModel：相同 API，更多性能

性能比较

Nemotron 3 Ultra 550B A55B（完整微调，多节点）

单节点 30B MoE 基准测试

Qwen3-30B-A3B

Nemotron 3 Nano 30B A3B

加速来源

Transformers v5 功能在 HuggingFace AutoModel 中的利用

专家后端

专家并行和 DeepEP

动态权重加载

入门指南

结论

Sources