< img id="wx_img" src="https://www.qbitai.com/wp-content/uploads/imgs/qbitai-logo-1.png" width="400" height="400">

英伟达MoE新开源:一行import,微调加速3.7倍

在Transformers v5的基础上,增加了专家并行、DeepEP和TransformerEngine

鱼羊 发自 凹非寺

量子位 | 公众号 QbitAI

一行import,MoE大模型微调提速3.7倍

英伟达最新研究成果现已开源:NeMo AutoModel,专为大规模构建和微调生成式AI模型而打造。

在Hugging Face Transformers v5的基础之上,NeMo AutoModel能做到不改代码API,只添一行import,就实现对MoE模型更快速的微调。

实验显示,相比Hugging Face原版Transformers v5,英伟达NeMo AutoModel能在MoE微调中实现3.4-3.7倍训练吞吐提升,并减少29%-32% GPU显存占用

在单节点8xH100 80GB GPU上,以Qwen3-30B-A3B为例,NeMo AutoModel直接把TPS/GPU(每GPU每秒吞吐量)从3075拉到11340,提升达到3.69倍。

核心技术解析

MoE已经成为当前前沿模型的主流架构,但MoE也给高效训练带来了新的挑战:

专家并行、通信融合、kernel优化……这些复杂工程都需要配套的基础设施来支持。

HuggingFace的Transformers v5是目前被用得比较多的MoE训练“通用底座”。v5增强了对MoE的原生支持,引入了expert backends、dynamic weight loading、分布式执行等MoE基础能力。

这一次,英伟达的思路就是站在前辈的肩膀上,兼容HuggingFace Transformers的API,让大家能不大改代码,就在MoE微调里获得更高训练吞吐和更低显存占用

具体来说,NeMo AutoModel在Transformers v5的基础上,增加了专家并行(EP)DeepEPTransformerEngine

专家并行(Expert Parallelism)

专家并行技术主要用来降低内存压力。

EP把专家权重分布到了多个GPU上,每张GPU不再完整持有所有expert,而是只持有其中一部分参数。

举个例子,8张GPU上ep_size=8,专家权重被分布至8块GPU,每张GPU的MoE内存占用能降到原来的1/8

从实验结果来看,对于Qwen3,这项技术能将峰值内存从68.2GiB降至48.1GiB,降幅29%。

对于Nemotron Nanomo模型,内存占用从62.1 GiB降至42.5 GiB,降幅 32%。

释放出的空间可以用来支持更大批次、更长的序列。

DeepEP

DeepEP实现了计算和通信的融合。

传统方式里,token分发和专家计算之间有明显通信成本。DeepEP把token分发和组合操作整合进优化的GPU内核,实现了通信过程和专家计算的重叠。

TransformerEngine

TransformerEngine内核为各类核心运算提供加速。

这项技术提供了融合注意力机制、线性层和RMSNorm等实现,不只加速MoE层,也加速普通Transformer层。

一行import,3倍速度提升

总结来说,对于原本就用上了Transformers v5的盆友们来说,英伟达NeMo AutoModel带来了一个无痛升级方案:

只需加上一行import代码,即可获得3倍MoE微调速度提升。

在Qwen3-30B-A3B和Nemotron 3 Nano 30B-A3B上,相较于Transformers v5,该方案可以实现3.4-3.7倍的训练吞吐量提升,同时内存消耗降低29%-32%。

英伟达还展示了Nemotron 3 Ultra 550B A55B在16个H100节点、128张GPU上的全参数微调结果。

TPS/GPU为815,TFLOP/s/GPU约为293,峰值内存为58.2GiB。

这里没跟v5对比的原因,是Transformers v5在这种规模下会直接撑爆内存¯_(ツ)_/¯

感兴趣的话,英伟达已经把代码、配置和基准测试脚本都放在GitHub上了:

https://github.com/NVIDIA-NeMo/Automodel/tree/blog/transformers-v5-automodel/blog_experiments

具体使用指南在这里:

https://docs.nvidia.com/nemo/automodel/latest/get-started/hf-compatibility

版权所有,未经授权不得以任何形式转载及使用,违者必究。