Accélération record en entraînement FP8 avec NVIDIA NeMo : performances et stabilité au cœur de la nouvelle génération d’IA

L’entraînement accéléré en précision FP8 avec NVIDIA NeMo s’inscrit dans une démarche visant à optimiser les performances réelles des modèles de langage à grande échelle (LLM) en exploitant pleinement les capacités matérielles des GPU NVIDIA Hopper et Blackwell. Ce post se concentre sur la vitesse d’entraînement — un enjeu critique en production — en comparant plusieurs recettes de mise à l’échelle FP8 via le cadre NVIDIA NeMo Framework 25.04. Les approches évaluées incluent le scaling actuel (sur H100), le scaling retardé, le scaling par sous-canals (sub-channel) et MXFP8 (sur DGX B200 et GB200). Les résultats montrent que l’efficacité FP8 dépend fortement de la granularité du scaling. Le scaling par tenseur (per-tensor) offre les meilleurs gains bruts de débit (throughput) dans les opérations GEMM, atteignant jusqu’à 2x par rapport au BF16, car il implique un seul facteur d’échelle par tenseur. Cependant, cette simplicité se paie par une stabilité numérique moindre, comme le montre une trajectoire de perte d’entraînement plus instable (ex : ligne verte claire sur la Figure 2), notamment pour les modèles complexes comme Llama 3.1 405B. En revanche, le scaling par blocs (block-wise), notamment MXFP8, applique un facteur d’échelle par blocs de 32 valeurs, ce qui améliore significativement la fidélité numérique. Bien que cela introduise un léger surcoût en gestion des échelles, il permet une convergence très proche du BF16 (ligne violette foncée), essentielle pour garantir la qualité du modèle final. Sur les GPU DGX B200, MXFP8 atteint un gain de 1,28 à 1,37x par rapport au BF16, avec une stabilité remarquable sur des modèles allant de 8B à 340B de paramètres. Ce comportement s’explique par l’optimisation matérielle du Blackwell : les Tensor Cores et la hiérarchie mémoire sont conçus pour gérer efficacement ces formats à petite échelle. Le superchip GB200 Grace Blackwell, combinant deux GPU B200 avec un CPU Grace via NVLink, offre un domaine mémoire partagé et une bande passante élevée. Il permet une accélération supplémentaire par rapport au B200 seul, particulièrement pour les modèles très gros (ex : Nemotron 340B), où les limites de communication deviennent prépondérantes. En résumé, le choix de la recette FP8 repose sur un compromis entre vitesse brute et stabilité numérique. Pour les modèles denses, plus ils sont grands, plus le gain FP8 est significatif — passant de 1,30x (8B) à 1,53x (405B) avec le scaling actuel sur H100. MXFP8, bien que légèrement moins performant en débit brut, offre une meilleure robustesse et est particulièrement adapté aux architectures Blackwell. Les experts du secteur soulignent que ces avancées rendent possible l’entraînement de modèles de 100 milliards de paramètres en quelques jours sur des infrastructures existantes, réduisant drastiquement les coûts énergétiques et financiers. NVIDIA NeMo, avec son intégration native du Transformer Engine, facilite l’adoption de ces techniques, offrant des configurations prêtes à l’emploi pour les chercheurs et ingénieurs. Ce progrès marque une étape clé vers une IA plus accessible, plus rapide et plus durable.