NVIDIA a renforcé son soutien aux algorithmes d'optimisation de pointe pour accélérer l'entraînement des grands modèles de langage (LLM), en mettant particulièrement l'accent sur le nouvel optimiseur Muon. Des méthodes d'optimisation de second ordre comme Shampoo sont utilisées depuis plus de dix ans, mais Muon, qui intègre la mise à jour orthogonalisée par l'itération de Newton-Schulz, a récemment permis l'entraînement de modèles open source de référence tels que Kimi K2 et GLM-5. NVIDIA détaille dans une nouvelle publication les technologies permettant de déployer ces optimiseurs à grande échelle sur son infrastructure. Des tests effectués sur le système NVIDIA GB300 NVL72 ont démontré l'efficacité de l'approche. En utilisant la bibliothèque NVIDIA NeMo Megatron Bridge 26.02, les mesures montrent une perte de performance marginale avec Muon par rapport à l'optimiseur AdamW standard. Pour le modèle Kimi K2 entraîné sur 256 GPU GB300, le débit de formation atteint 1 080 TFLOPs/s/GPU avec Muon, contre 1 051 pour AdamW. De même, pour le modèle Qwen3 30B-A3B sur huit GPU, Muon atteint 721 TFLOPs/s/GPU contre 713. L'utilisation des FLOPs (MFU) est même supérieure avec Muon lorsque l'on inclut les calculs supplémentaires des itérations de Newton-Schulz. L'adoption massive de ces optimiseurs de second ordre présentait plusieurs défis techniques. L'étape de préconditionnement, nécessitant souvent des décompositions complexes, augmentait considérablement les coûts de calcul et la consommation de mémoire, tout en risquant d'introduire une instabilité numérique. De plus, la distribution synchronisée des mises à jour sur des milliers de GPU pouvait créer des goulots d'étranglement de communication. NVIDIA a surmonté ces obstacles grâce à plusieurs innovations logicielles clés intégrées dans Megatron Core. Premièrement, l'entreprise a introduit un optimiseur distribué au niveau des couches. Contrairement aux approches traditionnelles qui répartissent les éléments d'une même couche entre plusieurs unités, cette méthode assigne chaque couche complète à un rang de parallélisme de données spécifique. Cela permet à chaque GPU de disposer de toutes les données nécessaires pour calculer les préconditionneurs complets, essentiels pour des optimiseurs comme Muon, sans nécessiter de communications excessives pour agréger des fragments de couches. Deuxièmement, une approche distribuée de l'itération de Newton-Schulz a été développée pour gérer les scénarios de parallélisme tensoriel. Trois modes de fonctionnement sont proposés : un mode dupliqué où toutes les données sont agrégées sur chaque GPU pour maximiser la vitesse réseau, un mode distribué qui répartit le calcul mathématique pour optimiser la puissance de traitement au prix de plus de communications, et un mode par blocs qui évite la communication mais ne calcule l'orthogonalisation que sur des sous-ensembles locaux. D'autres optimisations ont été mises en place pour améliorer l'efficacité globale, notamment le masquage des communications en chevauchant la récupération des paramètres avec le calcul du lot suivant, l'équilibrage de charge dynamique entre les couches de différentes tailles, et l'utilisation de noyaux SYRK et d'agrégation fusionnée pour réduire le volume de données transférées. Ces avancées, combinées au support d'autres optimiseurs émergents comme SOAP, permettent désormais à la communauté de recherche de déployer ces méthodes puissantes à l'échelle de milliers de GPU, repoussant ainsi les limites de l'efficacité d'entraînement des LLM.