使用NVIDIA NeMo AutoModel加速Transformer微调

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

MoE模型的兴起带来了新的训练效率挑战：跨数百个专家路由令牌、将专家矩阵乘法融合为单个内核、跨GPU分片权重以及重叠通信与计算，这些都需要通用库开箱即用之外的基础设施。Transformers v5引入了对MoE的一流支持，如专家后端、动态权重加载和分布式执行的张量并行计划。此外，v5通过将PyTorch的DeviceMesh直接集成到from_pretrained()中，使分布式训练成为一等公民。

NeMo AutoModel通过继承AutoModelForCausalLM，并添加专家并行、DeepEP融合全对全调度和TransformerEngine内核，在v5之上构建。DeepEP是v5尚未具备的部分：它将通信与专家计算重叠。由于NeMo AutoModel利用v5的可逆权重转换加载每个模型，它可以专注于这些可复用的核心操作，而非逐个模型的检查点适配，同时save_pretrained()仍输出标准HF检查点，供vLLM和SGLang等工具加载。

性能测试展示了NeMo AutoModel在两个场景下的优势：跨16节点全微调千亿级550B模型，以及单节点训练两个30B MoE模型。对于550B模型，Transformers v5因内存不足无法运行，而NeMo AutoModel通过专家并行（EP=64）将专家权重分片到各GPU，使全微调成为可能，达到每GPU每秒815个令牌，峰值内存58.2 GiB。在单节点8×H100 GPU上，Qwen3-30B-A3B模型从v5的3,075 TPS/GPU提升至11,340 TPS/GPU（3.69倍），峰值内存从68.2 GiB降至48.1 GiB（-29%）；Nemotron 3 Nano 30B A3B模型从4,583 TPS/GPU提升至15,421 TPS/GPU（3.36倍），峰值内存从62.1 GiB降至42.5 GiB（-32%）。

速度提升来自三个来源：专家并行降低内存压力（EP=8将每GPU MoE内存占用减少8倍）；DeepEP将通信与计算融合，避免单独的非阻塞通信；TransformerEngine内核加速核心操作，如融合注意力、线性层和RMSNorm。NeMo AutoModel还提供了可配置的后端，通过BackendConfig选择注意力、线性层、专家和调度器实现。总之，NeMo AutoModel为MoE模型微调提供了显著的性能提升，同时保持与HuggingFace Transformers的API兼容性。