NVIDIA fait passer le déterminisme bit à bit de Megatron Core à 2432 GPU avec seulement 2% de surcoût

Le problème de l'entraînement à grande échelle et la solution de NVIDIA

Lors de l'entraînement de modèles comptant des milliers de milliards de paramètres sur des milliers de GPU, les approches classiques de débogage s'effondrent. La multiplication des dimensions de parallélisme, les calculs à basse précision (MXFP8, NVFP4) et les points de contrôle distribués complexifient considérablement la reproduction des bugs et la validation des correctifs. NVIDIA a intégré dans Megatron Core le déterminisme bit à bit — une garantie que deux exécutions lancées depuis le même état initial suivront une trajectoire numérique strictement identique, à condition que les données, l'architecture, le parallélisme et les logiciels restent inchangés.

Megatron Core offre deux garanties clés : la reproductibilité indépendante (deux exécutions démarrées depuis le même état restent identiques bit à bit à chaque étape) et la reprise bit à bit depuis un point de contrôle (la restauration est identique à une exécution continue). Ce point est crucial pour l'entraînement de production, particulièrement coûteux — même un léger ralentissement, à l'échelle de milliers de milliards de paramètres, se traduit par des milliers de jours-GPU perdus.

Quatre bénéfices concrets

Le déterminisme bit à bit résout quatre problèmes liés au débogage et à la validation :

• Reproduction de bugs — figer un pic de loss à une étape précise, puis modifier un composant à la fois pour isoler la cause • Continuité des points de contrôle — reprendre un entraînement interrompu sans modifier la trajectoire numérique • Détection de défauts — exploiter les écarts entre des exécutions par ailleurs identiques pour repérer une corruption de point de contrôle ou des problèmes logiciels et matériels • Comparaison de l'impact des modifications — distinguer l'effet numérique d'un changement intentionnel de la variation naturelle entre exécutions

Méthodologie de débogage du non-déterminisme

NVIDIA a proposé, dans la PR #7262 de Megatron-LM, un système de traçage progressif enregistrant des empreintes ordonnées de tenseurs à travers les différents rangs. Le processus réduit progressivement le champ de recherche, des métriques de bout en bout jusqu'au kernel précis :

• Métriques de bout en bout — identifier la première itération où une métrique sérialisée diverge (loss d'entraînement, loss de modélisation du langage, norme du gradient, loss MTP) • Traçage sémantique large — couverture des opérations collectives, de la communication pipeline, du recomputation, des opérations de l'optimiseur et des réductions de gradient pour déterminer la phase d'entraînement divergente • Traçage au niveau des modules et des couches — localisation au niveau d'un composant du modèle (couche Transformer, bloc MoE, étape de l'optimiseur) • Traçage au niveau des opérations — empreintes des opérations ATen et des appels aux extensions pour identifier la première opération présentant des entrées identiques mais des sorties différentes • Outils au niveau des kernels — identification du kernel responsable et du mécanisme de non-déterminisme

Point crucial : il ne faut pas se limiter au traçage de l'itération où l'écart de loss devient visible, car le premier bit divergent peut apparaître bien avant. Le traçage large permet d'identifier l'itération et les rangs où la divergence apparaît le plus tôt ; le traçage détaillé n'est ensuite activé que sur cette fenêtre restreinte.

Les empreintes sont calculées directement sur GPU via torch.hash_tensor (avec enregistrement de la forme, du type de données et du nombre d'éléments). Pour MXFP8 et NVFP4, ce sont les valeurs encodées et les buffers d'échelle qui sont empreintés. Pour les données sensibles à l'ordre (routing maps, sorties de dispatch MoE), un fingerprinting par blocs (chunk-wise) est utilisé.

Optimisation pour le

Nemotron à mille milliards de paramètres

Une recette déterministe engendrant des surcoûts importants reste utile pour le débogage, mais s'avère trop coûteuse pour la production. NVIDIA a mené une co-optimisation des chemins déterministes et non déterministes sur le modèle hybride Nemotron, qui combine des couches SSM (state-space model) de type Mamba et des couches d'attention Transformer, en optimisant trois charges de travail :

• Nemotron 3 Ultra : réduction d'environ 17% à 2% après optimisation du remplissage des buffers, puis jusqu'à 1,5% sur 3072 GPU • Proxy hybride Triton : réduction d'environ 60% à 37% après restauration de la fusion MoE-MLP, puis jusqu'à 2% • Chemin de gradient de poids optimisé CuTeDSL : surcoût d'environ 3,6% sur 256 GPU, puis d'environ 2% sur 2432 GPU

Sur 2432 GPU, la recette de production a atteint un surcoût de déterminisme en régime stable d'environ 2%, tout en maintenant le déterminisme bit à bit pendant 800 étapes d'entraînement. Les optimisations au niveau des kernels comprennent l'allocation d'emplacements de sortie privés pour les writers du grouped-GEMM, ainsi que l'application d'un ordre de réduction fixe pour préserver le parallélisme tout en garantissant des résultats déterministes. Megatron Core valide les recettes, teste les kernels et vérifie les modules pour l'ensemble des configurations de parallélisme, y compris la prise en charge des formats à basse précision FP8 et FP4.