使用 NVIDIA NeMo AutoModel 加速 Transformers 微调

来源：https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel 返回文章列表 (https://huggingface.co/blog)

• NeMo AutoModel：相同API，更高性能 (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#nemo-automodel-same-api-more-performance)
• 性能对比 (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#performance-comparison)
  • Nemotron 3 Ultra 550B A55B（完整微调，多节点）(https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#nemotron-3-ultra-550b-a55b-full-fine-tune-multi-node)
  • 单节点30B MoE基准测试 (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#single-node-30b-moe-benchmarks)
    • Qwen3-30B-A3B (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#qwen3-30b-a3b)
    • Nemotron 3 Nano 30B A3B (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#nemotron-3-nano-30b-a3b)
    • 加速的来源 (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#where-the-speedup-comes-from)
• Transformers v5中被HuggingFace AutoModel利用的特性 (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#transformers-v5-features-leveraged-by-huggingface-automodel)
  • 专家后端 (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#expert-backends)
• 专家并行与DeepEP (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#expert-parallelism-and-deepep)
  • 动态权重加载 (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#dynamic-weight-loading)
• 快速上手 (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#getting-started)
• 结论 (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#conclusion)
• 致谢 (https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#acknowledgements)

HuggingFace Transformers已成为开源AI生态系统的基础，最近的Transformers v5 (https://github.com/huggingface/transformers/releases/tag/v5.0.0) 版本通过为混合专家（MoE）模型提供一流支持而进一步巩固了其地位，而MoE模型现在已成为前沿模型 (https://www.nvidia.com/en-us/glossary/frontier-models/) 的主导架构。v5提供了MoE的基础设施：专家后端、动态权重加载以及分布式执行，使得MoE具有可扩展性且易于构建。

NVIDIA NeMo AutoModel (https://github.com/NVIDIA-NeMo/Automodel) 是一个开源库，隶属于NVIDIA NeMo框架 (https://github.com/NVIDIA-NeMo)，用于大规模构建自定义生成式AI模型。NeMo AutoModel 在v5的基础上进行了简洁的构建，增加了专家并行、DeepEP融合全对全分发以及TransformerEngine内核，并利用v5的动态权重加载将这些优化应用于广泛且不断增长的模型家族。其回报是：在使用相同的 from_pretrained() API（仅需一行导入，无需其他代码更改）微调MoE模型时，训练吞吐量提高 3.4-3.7倍，GPU内存减少 29-32%。

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

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#background 背景

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

Transformers v5 (https://github.com/huggingface/transformers/releases/tag/v5.0.0)（简称“v5”）引入了对MoE的一流支持，例如专家后端 (https://huggingface.co/docs/transformers/en/experts_interface)、动态权重加载 (https://huggingface.co/docs/transformers/en/weightconverter) 以及用于分布式执行的张量并行方案。此外，v5通过将PyTorch的 DeviceMesh 直接集成到 from_pretrained() 中，使分布式训练成为一等公民。

NeMo AutoModel (https://github.com/NVIDIA-NeMo/Automodel) 在v5基础上构建，通过继承 AutoModelForCausalLM，并增加了专家并行（EP）、DeepEP融合全对全分发以及TransformerEngine内核。DeepEP是v5尚未具备的部分：它将通信与专家计算重叠。由于NeMo AutoModel利用v5的可逆权重转换来加载每个模型，它可以将其工程精力集中在这些可重用的核心操作上，而不是特定模型的检查点管道中，而 save_pretrained() 仍然生成标准的HF检查点，可以被vLLM和SGLang等工具加载。

下一节将介绍两者如何协同工作，以及我们测量到的性能提升——从在16个节点上完整微调NVIDIA Nemotron 3 Ultra 550B A55B (https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16)，到单节点模型如Qwen3-30B-A3B和Nemotron 3 Nano 30B A3B (https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16)。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#nemo-automodel-same-api-more-performance NeMo AutoModel：相同API，更高性能

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

以下是在两者中加载模型的方式。唯一改变的是导入语句：

nemo_and_hf (https://cdn-uploads.huggingface.co/production/uploads/690d0a6c2c5acfe0e1f4777d/VTPq2Wp-RrEcP1eGUJxao.png)

这一行导入完成了大量工作。对于流行的MoE架构，如Qwen3、NVIDIA Nemotron (https://developer.nvidia.com/nemotron)、GPT-OSS和DeepSeek V3，NeMo AutoModel提供了手工调优的实现 (https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/_transformers/registry.py)，采用TransformerEngine注意力、融合线性层和自定义专家内核。对于其他架构，它回退到标准的HF实现，同时仍然应用Liger内核 (https://github.com/linkedin/Liger-Kernel) 补丁等优化。无论采用哪种路径，生成的模型都随时可以扩展：传递一个 device_mesh，即可实现多GPU训练，无需进一步重写。

NeMo AutoModel真正闪耀之处在于将MoE模型扩展到多GPU训练。要使用专家并行在8个GPU上训练Nemotron 3 Nano 30B A3B (https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16)，只需添加分布式网格配置：

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分发）。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#performance-comparison 性能对比

我们在两种场景下评估了NeMo AutoModel：在16个节点上完整微调前沿规模的550B模型，以及在单个节点上训练两个30B MoE模型。550B的结果展示了专家并行在大规模场景下的必要性；30B的结果则量化了相对于Transformers v5的每GPU加速比。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#nemotron-3-ultra-550b-a55b-full-fine-tune-multi-node Nemotron 3 Ultra 550B A55B（完整微调，多节点）

Nemotron 3 Ultra 550B A55B (https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16) 是一个550B参数的混合模型，集成了Mamba2、LatentMoE和多令牌预测（MTP）。我们进行了 完整微调：每个参数都更新，并实例化了Adam优化器状态，在这个规模下需要 16个H100节点（128个GPU）。

方法：

为何没有Transformers v5列。 在此规模下，Transformers v5内存不足，因此没有v5的数字可报告。AutoModel的专家并行将专家分片到多个GPU上，使内存占用保持在预算范围内，从而使得完整微调得以运行。下面的30B对比显示了v5同样适用时的相同优势。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#single-node-30b-moe-benchmarks 单节点30B MoE基准测试

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

方法：

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

nemo_automodel_blog_chart_mockup_v5 (https://cdn-uploads.huggingface.co/production/uploads/690d0a6c2c5acfe0e1f4777d/rbCVgV6a18c4UcDsiWfZN.png)

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#qwen3-30b-a3b Qwen3-30B-A3B

v4死锁原因： Transformers v4将Qwen3 MoE专家存储为包含128个独立MLP模块的ModuleList，每个模块单独进行FSDP包装。前向传播使用数据依赖循环，仅迭代接收到令牌的专家。由于不同rank上的数据不同，不同rank会跳过不同的专家，导致FSDP AllGather/ReduceScatter集合不匹配，从而无限挂起。Transformers v5通过将专家存储为融合的3D参数张量（没有按专家的模块，没有按专家的FSDP集合）来修复此问题。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#nemotron-3-nano-30b-a3b Nemotron 3 Nano 30B A3B

v4配置： trust_remote_code=True（NVIDIA的hub建模代码）。hub代码的专家循环是FSDP安全的（遍历所有专家，无论令牌分配如何），因此不会像Qwen3 v4那样死锁。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#where-the-speedup-comes-from 加速的来源

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集合，而是将令牌分发和结果组合融合到优化的GPU内核中，将通信与专家计算重叠。
3. TransformerEngine内核加速核心操作。 TE的融合注意力、线性层和RMSNorm实现相对于其PyTorch/Flash Attention等效方案，在所有层类型（不仅仅是MoE层）上都提供了一致的加速。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#transformers-v5-features-leveraged-by-huggingface-automodel Transformers v5中被HuggingFace AutoModel利用的特性

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#expert-backends 专家后端

Transformers v5中最有影响力的特性之一是 experts_implementation (https://huggingface.co/docs/transformers/en/experts_interface) 参数，它包括三个专家后端：

grouped_mm 后端是关键的训练优化：它不是逐个循环专家，而是根据分配的专家对令牌进行排序，并执行单一的融合分组矩阵乘法。

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",           # TransformerEngine注意力
    linear="te",         # TransformerEngine线性层
    experts="torch_mm",  # 分组专家矩阵乘法
    dispatcher="deepep", # DeepEP融合全对全
)

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#expert-parallelism-and-deepep 专家并行与DeepEP

Transformers v5也提供了一个专家并行路径 (https://huggingface.co/docs/transformers/en/expert_parallelism)。它将专家权重分片到多个GPU上。GroupedGemmParallel (https://github.com/huggingface/transformers/blob/v5.10.2/src/transformers/integrations/tensor_parallel.py#L1078) 风格仅加载每个设备的本地专家，而 RouterParallel (https://github.com/huggingface/transformers/blob/v5.10.2/src/transformers/integrations/tensor_parallel.py#L1123) 则路由令牌并通过 all_reduce 结合结果。它巧妙地构建在v5现有的张量并行机制之上。启用它后，模型会使用 from_pretrained() 自动加载。NeMo AutoModel通过DeepEP (https://github.com/deepseek-ai/DeepEP) 扩展了这一点——一个用于专家通信的GPU内核库，它使用一种称为“融合全对全”的技术来将令牌分发与专家计算重叠。这种通信-计算重叠是NeMo AutoModel加速的主要来源之一，并且在v5本身尚未实现。

在NeMo AutoModel内部，专家并行是通过修改模型的 forward 和 backward 来使用DeepEP的 dispatch 和 combine 操作而实现的。这发生在底层的自定义MoE实现中，但对于最终用户是完全透明的：只需指定 ep_size 即可激活。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#dynamic-weight-loading 动态权重加载

NeMo AutoModel利用v5的动态权重加载 (https://huggingface.co/docs/transformers/en/weightconverter) 系统，其中一个定义 convert.py 和 registry.py 的社区贡献模型能够从标准HF检查点恢复，并自动应用自定义优化。这使得NeMo AutoModel可以支持任何在HF Hub上有的模型：模型发布者不必做任何额外工作；社区贡献者只需定义权重转换即可。

结果是一个可扩展的生态：无需调整即可支持成千上万个模型，同时优化几千个可能需要手动实现的可选的模型。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#getting-started 快速上手

开始使用NeMo AutoModel。

1. 安装

pip install nemo-automodel

这从PyPI安装了NeMo AutoModel，包括所有依赖。

2. 微调Qwen3-30B-A3B（单节点，8个GPU）

import os
import torch
from torch.distributed import init_process_group, destroy_process_group
from nemo_automodel import NeMoAutoModelForCausalLM
from nemo_automodel.recipes._dist_utils import create_distributed_setup_from_config
from nemo_automodel.recipes.utils import train

init_process_group(backend="nccl")
torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

dist_setup = create_distributed_setup_from_config({
    "strategy": "fsdp2",
    "ep_size": 8,
})

model = NeMoAutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-30B-A3B",
    dtype=torch.bfloat16,
    distributed_setup=dist_setup,
)

tokenizer = model.get_tokenizer()

# train() 函数封装了前向、损失和反向
train(model, tokenizer, ...)  # 插入训练循环

destroy_process_group()

3. 微调Nemotron 3 Ultra 550B A55B（多节点，128个GPU） 在多节点设置中，工作流程类似，但通过 distributed_setup 传入了更大的EP配置。模型加载会处理不同节点间的专家分片：

dist_setup = create_distributed_setup_from_config({
    "strategy": "fsdp2",
    "ep_size": 64,  # 跨128个GPU分片
})

model = NeMoAutoModelForCausalLM.from_pretrained(
    "nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16",
    dtype=torch.bfloat16,
    distributed_setup=dist_setup,
)

有关完整指南，还包括多节点设置、多数据集训练、日志记录和评估的详细配方，请参阅NeMo AutoModel文档 (https://github.com/NVIDIA-NeMo/Automodel) 和示例 (https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples)。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#conclusion 结论

Transformers v5为MoE训练提供了强大的基础——专家后端、动态权重加载和分布式执行——使其成为整个开源AI生态系统的基石。NeMo AutoModel在此基础上构建，通过DeepEP将通信与计算重叠、用TransformerEngine替换核心操作，并支持专家并行以减少MoE内存占用，实现了3.4-3.7倍的训练吞吐量提升和29-32%的内存减少。

结果是一个统一的微调路径，在NeMo AutoModel的 from_pretrained() 和 save_pretrained() API之后，可以透明地处理从单GPU探索到多节点训练的一切。对于现代MoE模型（Qwen3、DeepSeek V3、Nemotron）和标准密集模型，NeMo AutoModel提供了一种简单、高性能的微调方式——无需更改API。

https://huggingface.co/blog/nvidia/accelerating-fine-tuning-nvidia-nemo-automodel#acknowledgements 致谢

作者感谢HuggingFace Transformers团队提供的扎实的v5基础、DeepEP团队提供的专家通信库，以及整个NVIDIA NeMo团队提供的基础支持。