NVIDIA accélère le préentraînement des LLM avec NVFP4 sur Blackwell, JAX et MaxText

Le préentraînement des LLM frontier se heurte à la limite de débit des systèmes de calcul. NVIDIA a montré comment l'association de JAX, MaxText et du nouveau format NVFP4 sur les puces Blackwell permet d'accélérer nettement ce processus sans perte de qualité.

Pourquoi chaque point de pourcentage compte

Lorsque l'entraînement porte sur des milliers de milliards de tokens répartis sur des milliers d'accélérateurs, gagner ne serait-ce qu'un pour cent de temps à chaque étape représente plusieurs jours de temps réel. À l'échelle frontier, cela équivaut directement à des millions de dollars de coûts de calcul.

NVFP4, un format à virgule flottante de quatre bits apparu avec l'architecture Blackwell, est devenu l'un des outils clés pour accélérer les opérations matricielles. Par rapport au FP8, il compacte les nombres deux fois plus densément, ce qui allège la charge mémoire et augmente le débit effectif des cœurs tensoriels.

La principale difficulté : la grille numérique à quatre bits est clairsemée. Mal configurés, les gradients sortent facilement de sa plage, ce qui fait diverger l'entraînement. NVIDIA et l'équipe de MaxText ont résolu ce problème avec des schémas de mise à l'échelle sur mesure et un loss scaling dynamique.

Comment fonctionne la précision mixte avec NVFP4

L'entraînement en précision mixte n'est pas une approche nouvelle : FP8 et BF16 sont déjà des standards de l'industrie. NVFP4 va plus loin en autorisant des poids en 4 bits dans les multiplications matricielles les plus lourdes, tout en conservant une précision supérieure là où elle compte vraiment.

• NVFP4 s'applique aux poids et aux activations dans les opérations GEMM
• BF16 ou FP32 restent utilisés pour les accumulateurs et la normalisation
• MaxText oriente automatiquement les opérations vers le bon format
• JAX compile le graphe de calcul via XLA, en optimisant les kernels pour Blackwell
• Résultat : un throughput accru pour une consommation énergétique comparable ou inférieure

La pile logicielle et ce qu'il faut changer dans le code

MaxText est un framework d'entraînement open source haute performance reposant sur JAX et développé par Google. Conçu à l'origine pour les TPU, il est activement adapté aux clusters de GPU, et le partenariat avec NVIDIA n'a rien de surprenant.

NVIDIA a intégré des kernels NVFP4 de bas niveau à cuBLAS et cuDNN, et JAX/XLA prend en charge ces opérations via des adaptateurs dédiés. Les développeurs n'ont pas à réécrire le code d'entraînement à la main : il suffit d'activer les bons flags dans la configuration de MaxText et de vérifier que le cluster est équipé de puces Blackwell (B100, B200, GB200).

«

La précision numérique est l'un des paramètres au plus fort effet de levier, mais le préentraînement en précision mixte à faible nombre de bits est difficile à réussir », note l'équipe du NVIDIA Developer Blog.

Ce que cela change

Pour les équipes qui préentraînent des modèles frontier, NVFP4 sur Blackwell représente une accélération quasi gratuite : la pile JAX et MaxText existante n'exige que des modifications minimales de configuration. À l'échelle de centaines ou de milliers de GPU, même un gain de 10-15 % de throughput réduit directement le time-to-checkpoint et le budget de calcul global.

La course à l'efficacité du préentraînement entre dans une phase de bataille pour la précision numérique.