NVIDIA NeMo AutoModel : Révolutionner le fine-tuning des Transformers

L'essor des Transformers et l'innovation de NVIDIA

Dans le domaine de l'intelligence artificielle, les Transformers de HuggingFace se sont imposés comme une pierre angulaire de l'écosystème open-source. La sortie de la version Transformers v5 a consolidé cette position, notamment grâce à l'intégration des modèles Mixture-of-Experts (MoE), qui sont devenus le standard pour les architectures de pointe. Cette version a introduit des fonctionnalités essentielles telles que les backends d'experts, le chargement dynamique de poids, et l'exécution distribuée, rendant les MoE plus extensibles et plus simples à construire.

Dans ce contexte, NVIDIA NeMo AutoModel se distingue comme une bibliothèque ouverte intégrée au cadre NVIDIA NeMo, dédiée à la création de modèles d'IA générative personnalisés à grande échelle. En s'appuyant sur les innovations de v5, NeMo AutoModel introduit des améliorations significatives, notamment le parallélisme d'experts, le dispatch all-to-all fusionné DeepEP, et les noyaux TransformerEngine. Ces optimisations permettent un débit d'entraînement 3,4 à 3,7 fois supérieur et une réduction de 29 à 32 % de la mémoire GPU lors du fine-tuning des modèles MoE par rapport aux Transformers v5 natifs, tout en utilisant la même API from_pretrained().

Ce blog détaille comment cette combinaison fonctionne et comment les utilisateurs peuvent affiner les modèles MoE plus rapidement sans changer leurs API.

Défis et solutions pour l'entraînement des modèles MoE

L'adoption croissante des modèles MoE a introduit de nouveaux défis pour l'entraînement efficace des modèles d'IA. Parmi ces défis figurent le routage des tokens à travers des centaines d'experts, la fusion des multiplications de matrices d'experts en un seul noyau, le partitionnement des poids à travers les GPU, et le chevauchement de la communication avec le calcul. Ces exigences nécessitent une infrastructure bien au-delà de ce qu'une bibliothèque généraliste peut offrir par défaut.

La version Transformers v5 a apporté un support de premier ordre pour les MoE, en introduisant des backends d'experts, le chargement dynamique de poids, et des plans de parallélisme tensoriel pour l'exécution distribuée. De plus, v5 a intégré le DeviceMesh de PyTorch directement dans la fonction from_pretrained(), facilitant ainsi l'entraînement distribué.

NeMo AutoModel s'appuie sur ces fondations en sous-classant AutoModelForCausalLM et en ajoutant des fonctionnalités comme le parallélisme d'experts (EP), le dispatch all-to-all fusionné DeepEP, et les noyaux TransformerEngine. DeepEP, en particulier, chevauche la communication avec le calcul des experts, optimisant ainsi l'efficacité de l'entraînement. Grâce à la conversion de poids réversible de v5, NeMo AutoModel peut se concentrer sur ces opérations de base réutilisables, tout en permettant à save_pretrained() de produire des checkpoints compatibles avec des outils comme vLLM et SGLang.

La section suivante explique comment les deux fonctionnent ensemble et les gains de performance que nous avons mesurés, allant du fine-tuning complet du NVIDIA Nemotron 3 Ultra 550B A55B sur 16 nœuds jusqu'à des modèles à nœud unique comme Qwen3-30B-A3B et Nemotron 3 Nano 30B A3B.

Compatibilité API et performance accrue

Un des principaux objectifs de NeMo AutoModel est de maintenir la compatibilité avec l'API de HuggingFace Transformers, permettant ainsi à la communauté open-source de l'adopter facilement. En sous-classant AutoModelForCausalLM, NeMo AutoModel garantit que tout code fonctionnant avec les modèles HF est également compatible avec AutoModel.

Le chargement d'un modèle avec NeMo AutoModel est simple et ne nécessite qu'une modification mineure de l'importation. Cette seule importation active une série d'optimisations pour des architectures MoE populaires telles que Qwen3, NVIDIA Nemotron, GPT-OSS, et DeepSeek V3. NeMo AutoModel fournit des implémentations optimisées avec TransformerEngine attention, des couches linéaires fusionnées, et des noyaux d'experts personnalisés. Pour les autres modèles, il revient aux standards HF tout en appliquant des optimisations comme le patching de noyau Liger.

Là où NeMo AutoModel excelle vraiment, c'est dans la mise à l'échelle des modèles MoE pour un entraînement multi-GPU. Par exemple, pour entraîner le modèle Nemotron 3 Nano 30B A3B avec le parallélisme d'experts sur 8 GPU, il suffit d'ajouter une configuration de maillage distribué spécifique. Cela offre rapidité, évolutivité et optimisations de mémoire avec FSDP2, le parallélisme d'experts, les noyaux TransformerEngine, et le dispatch DeepEP, le tout à partir d'un simple appel from_pretrained().

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", model = NeMoAutoModelForCausalLM.from_pretrained( "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", dtype=torch.bfloat16, distributed_setup=dist_setup, ) )

dist.destroy_process_group()

## Analyse comparative des performances

Pour évaluer l'efficacité de NeMo AutoModel, deux scénarios ont été testés : le fine-tuning complet d'un modèle de 550 milliards de paramètres sur 16 nœuds, et l'entraînement de deux modèles MoE de 30 milliards de paramètres sur un seul nœud. Le modèle **Nemotron 3 Ultra 550B A55B**, un hybride de 550 milliards de paramètres, a été utilisé pour le benchmark de fine-tuning complet. Ce modèle, équipé de **Mamba2**, **LatentMoE**, et **Multi-Token Prediction (MTP)**, a été testé sur 16 nœuds H100 (128 GPU).

Dans ce test, le parallélisme d'experts a permis de répartir les experts à travers les GPU, réduisant ainsi l'empreinte mémoire et permettant au fine-tuning complet de s'exécuter efficacement. Les tests sur des modèles de 30 milliards de paramètres ont montré des gains similaires, où NeMo AutoModel a surpassé Transformers v5 en termes de débit et d'efficacité mémoire.

### Nemotron 3 Ultra 550B A55B (fine-tuning complet, multi-nœud)

Le modèle **Nemotron 3 Ultra 550B A55B** est un modèle hybride de 550 milliards de paramètres livré avec **Mamba2**, **LatentMoE**, et **Multi-Token Prediction (MTP)**. Nous avons effectué un benchmark de fine-tuning complet : chaque paramètre est mis à jour et l'état de l'optimiseur Adam est matérialisé, ce qui à cette échelle s'étend sur 16 nœuds H100 (128 GPU).

- 16x H100 80GB (128 GPU)
- Parallélisme d'experts
- Taille de lot local
- MTP, activation checkpointing, entropie croisée linéaire fusionnée
- Dispatch DeepEP + experts torch_mm + TransformerEngine

Pourquoi il n'y a pas de colonne pour Transformers v5 ? Transformers v5 manque de mémoire à cette échelle, donc il n'y a pas de chiffre v5 à rapporter ici. Le parallélisme d'experts d'AutoModel répartit les experts à travers les GPU pour ramener l'empreinte dans les limites, ce qui permet au fine-tuning complet de s'exécuter. Les comparaisons de 30 milliards ci-dessous montrent le même avantage où v5 s'adapte.

### Benchmarks MoE de 30 milliards à nœud unique

Nous avons benchmarké trois approches sur un seul nœud avec 8x H100 80GB GPU : **HF Transformers v4** (code hub), **HF Transformers v5** (avec les meilleures optimisations disponibles), et **NeMo AutoModel** (EP=8 + noyaux personnalisés).

- 8x H100 80GB (nœud unique)
- Taille de lot local

Une note sur la porte de routage. Les chiffres de NeMo AutoModel ci-dessous utilisent une porte de routage équilibrée, qui force les tokens à être distribués uniformément entre les experts. Cela émule le point de fonctionnement idéal vers lequel un MoE est entraîné : la perte d'équilibrage de charge d'un modèle bien entraîné pousse l'utilisation des experts vers une uniformité presque parfaite, donc le routage équilibré reflète l'état stable vers lequel une charge de travail réelle converge (et élimine le bruit des stragglers que des tokens aléatoires pourraient injecter dans le parallélisme d'experts). Les colonnes v4/v5 exécutent leur routeur natif sur les mêmes tokens fictifs. La porte équilibrée mesure donc NeMo AutoModel à son point de fonctionnement MoE cible, et les colonnes v4/v5 reflètent leur comportement par défaut.

## Pourquoi l'accélération provient de là

L'accélération de **3,4 à 3,7 fois** de NeMo AutoModel par rapport à Transformers v5 provient de trois sources :

- **Le parallélisme d'experts** réduit la pression mémoire. EP=8 répartit les poids des experts à travers les GPU, réduisant l'empreinte MoE par GPU de 8x. Pour Qwen3, cela fait chuter la mémoire maximale de **68,2 GiB** à **48,1 GiB** (-29 %). Pour Nemotron Nano, cela passe de **62,1 GiB** à **42,5 GiB** (-32 %), libérant ainsi de la marge pour des tailles de lot plus grandes ou des séquences plus longues.

- **DeepEP fusionne communication et calcul**. Au lieu de collectes AllGather/ReduceScatter séparées pour le routage des experts, DeepEP fusionne le dispatch des tokens et combine en noyaux GPU optimisés, chevauchant la communication avec le calcul des experts.

- **Les noyaux TransformerEngine** accélèrent les opérations de base. Les implémentations d'attention fusionnée, de couches linéaires et de **RMSNorm** de TE offrent des accélérations constantes par rapport à leurs équivalents PyTorch/Flash Attention à travers tous les types de couches, pas seulement les couches MoE.

## Fonctionnalités de Transformers v5 Exploitées par HuggingFace AutoModel

Une des fonctionnalités les plus impactantes de Transformers v5 est le paramètre **experts_implementation**, qui inclut trois backends d'experts :

- Boucle for sur les experts sélectionnés
- Débogage, compatibilité et correction. Également disponible pour v4.
- Duplique les paramètres d'experts, GEMM groupé unique via **torch.bmm**
- Petits inputs, rapide avec **torch.compile**. Ajouté pour v5.
- Trie les tokens par expert, GEMM groupé unique via **torch.nn.functional.grouped_mm**
- Entraînement (économe en mémoire, pas de duplication de paramètres). Ajouté pour v5.

Le backend **grouped_mm** est l'optimisation clé pour l'entraînement : au lieu de boucler sur les experts un par un, il trie les tokens par leur expert assigné et exécute une seule multiplication de matrices groupée fusionnée.

NeMo AutoModel va plus loin. Pour les modèles avec des implémentations personnalisées, il utilise le dispatch all-to-all fusionné **DeepEP** combiné avec des noyaux **GMM** groupés et des couches linéaires **TransformerEngine**. La progression se présente comme suit :

- v4 (boucle for impatiente) → v5 (grouped_mm) → NeMo AutoModel (DeepEP + GMM + TE)

Dans NeMo AutoModel, le backend d'expert est configuré via **BackendConfig** :

```python
from nemo_automodel.components.models.common.utils import BackendConfig

backend = BackendConfig(
    experts="torch_mm",
    dispatcher="deepep",
)

En résumé, NVIDIA NeMo AutoModel représente une avancée majeure dans le fine-tuning des modèles Transformers, offrant des performances accrues et une compatibilité API qui facilitent son adoption par la communauté IA.