NeMo AutoModel de NVIDIA accélère 3,7 fois le fine-tuning des modèles MoE grâce au parallélisme expert et DeepEP

NVIDIA a déployé NeMo AutoModel, une bibliothèque open source au sein du framework NeMo qui apporte des gains de performance significatifs au fine-tuning des modèles mixture-of-experts (MoE). Construite directement sur Hugging Face Transformers v5, la bibliothèque offre un débit d'entraînement 3,4 à 3,7 fois supérieur et réduit la consommation de mémoire GPU de 29 à 32 % par rapport à la meilleure configuration disponible de Transformers v5, tout en conservant la même API from_pretrained(). Les utilisateurs n'ont qu'à changer une seule ligne d'importation pour bénéficier de ces avantages.

Contexte : Le défi de l'efficacité des MoE

L'essor des modèles MoE a introduit de nouveaux obstacles pour un entraînement efficace. Acheminer des jetons vers des centaines d'experts, fusionner les multiplications matricielles experts en un seul noyau, répartir les poids entre les GPU et chevaucher la communication avec le calcul nécessitent une infrastructure que les bibliothèques à usage général ne fournissent généralement pas par défaut.

Transformers v5 a ajouté un support MoE de première classe, y compris des backends experts, un chargement dynamique des poids et des plans de parallélisme tensoriel pour l'exécution distribuée. Mais des éléments clés de performance, comme le chevauchement efficace de la communication, sont restés non traités. NeMo AutoModel comble cette lacune.

Comment fonctionne NeMo AutoModel

NeMo AutoModel sous-classe AutoModelForCausalLM et intègre trois optimisations majeures : le parallélisme expert (EP), le dispatch fusionné tout-à-tout de DeepEP et les noyaux TransformerEngine. DeepEP chevauche la communication avec le calcul expert, une capacité que Transformers v5 ne possède pas. Comme NeMo AutoModel exploite la conversion réversible des poids de v5 pour charger les modèles, il concentre l'effort d'ingénierie sur des opérations de base réutilisables plutôt que sur la plomberie de points de contrôle par modèle.

La bibliothèque reste entièrement compatible avec l'API Hugging Face Transformers. Charger un modèle ne nécessite qu'un changement dans la déclaration d'importation, comme le montrent les exemples de code de NVIDIA.

Benchmarks de performance

Multi-nœuds : Nemotron 3 Ultra 550B A55B

Pour le fine-tuning complet du modèle Nemotron 3 Ultra A55B de 550 milliards de paramètres sur 16 nœuds H100 (128 GPU), NeMo AutoModel avec un parallélisme expert fixé à 64 a atteint 815 jetons par seconde par GPU, 293 TFLOP par seconde par GPU et une mémoire maximale de 58,2 GiB. De manière cruciale, Transformers v5 manque de mémoire à cette échelle, rendant une comparaison directe impossible sans EP.

Nœud unique : Qwen3-30B-A3B

Sur un seul nœud avec 8 GPU H100 de 80 Go, NeMo AutoModel (EP=8) a enregistré 11 340 TPS/GPU, soit un bond de 3,69 fois par rapport aux 3 075 TPS/GPU de Transformers v5. La mémoire maximale est passée de 68,2 GiB à 48,1 GiB, soit une réduction de 29 %. Notamment, Transformers v4 a complètement bloqué sur ce modèle en raison de l'encapsulation FSDP incompatible des modules experts individuels.

Nœud unique : Nemotron 3 Nano 30B A3B

Pour le modèle Nemotron Nano de 30 milliards, NeMo AutoModel a atteint 15 421 TPS/GPU, soit une amélioration de 3,36 fois par rapport aux 4 583 TPS/GPU de v5. La consommation maximale de mémoire est passée de 62,1 GiB à 42,5 GiB, soit une réduction de 32 %.

Sources de l'accélération

Les gains de performance proviennent de trois optimisations coordonnées :

• Le parallélisme expert répartit les poids experts entre les GPU. Avec EP=8 sur 8 GPU, chaque GPU ne détient qu'un huitième des paramètres experts. Pour un modèle avec environ 55 GiB de poids experts, EP réduit l'empreinte par GPU à environ 6,8 GiB, permettant un entraînement là où les approches FSDP seules manquent de mémoire.
• DeepEP fusionne l'envoi et la combinaison des jetons en noyaux GPU optimisés, chevauchant la communication avec le calcul expert plutôt que d'utiliser des collectifs AllGather et ReduceScatter séparés.
• Les noyaux TransformerEngine accélèrent les opérations de base comme l'attention fusionnée, les couches linéaires et RMSNorm, offrant des accélérations constantes sur tous les types de couches.

Exploitation des fonctionnalités de Transformers v5

NeMo AutoModel s'appuie sur plusieurs fonctionnalités clés de v5, notamment le système de backends experts. v5 a introduit le paramètre experts_implementation avec trois options : eager (débogage par boucle), batched_mm (GEMM unique par lot) et grouped_mm (exécution triée des jetons). NeMo AutoModel étend cela en utilisant le dispatch fusionné tout-à-tout de DeepEP combiné à des noyaux grouped GEMM et des couches linéaires TransformerEngine.

Le chargement dynamique des poids via WeightConverter et WeightRenaming de v5 permet aux points de contrôle MoE d'être stockés sous forme de tenseurs 3D fusionnés. Plus de 20 types de modèles utilisent ce mécanisme, notamment Mixtral, Qwen2 MoE, Qwen3 MoE, DeepSeek V2/V3 et OLMoE. Les conversions sont entièrement réversibles : save_pretrained() produit des tensors de sécurité au format Hugging Face standard que des outils en aval comme vLLM et SGLang peuvent charger.

Conclusion et disponibilité

NeMo AutoModel de NVIDIA offre une voie de mise à niveau sans friction pour les utilisateurs de Hugging Face qui souhaitent passer à l'échelle dans l'entraînement de modèles. En changeant une seule ligne d'importation, les utilisateurs obtiennent une instance de modèle plus de trois fois plus rapide avec une empreinte mémoire bien plus petite. Le code, les configurations et les scripts de benchmark sont disponibles dans le dépôt NeMo AutoModel. Les contributeurs principaux de ce travail incluent Adil Asif, Hemil Desai, Alexandros Koumparoulis et Huiying Li.