NVIDIA teste les méthodes FP32 sur Tensor Cores avec PET; les gains des noyaux et des pas complets diffèrent

NVIDIA a publié le 29 octobre une étude de précision et de vitesse pour les potentiels interatomiques appris par machine. Elle figure dans notre édition technique du 1er novembre, sans modifier la date originale. L’étude porte sur des charges précises du Point Edge Transformer, ou PET, et sur des GPU déterminés. Elle ne prouve pas que tous les modèles scientifiques obtiennent la même accélération. Il faut distinguer la vitesse d’un noyau de multiplication matricielle de celle d’un pas complet de dynamique moléculaire.

Un potentiel interatomique prédit l’énergie d’une configuration d’atomes; les forces dérivées de cette énergie guident ensuite les positions et vitesses sur de nombreuses étapes. De petites erreurs numériques peuvent s’accumuler dans une trajectoire, d’où le maintien de FP32 dans beaucoup de modèles. La méthode de NVIDIA ne réentraîne pas PET et ne change pas son architecture. Elle sépare les opérandes FP32 en composants de précision inférieure, calcule certains produits sur Tensor Cores et accumule les résultats en FP32. BF16x9 sur Blackwell utilise trois composants par opérande et neuf produits. TF32x3 sur Ampere retient trois produits et omet le plus petit produit entre résidus. Leurs profils de précision et de vitesse restent différents.

Les essais précisent leurs dénominateurs. NVIDIA emploie un petit modèle PET-MAD de 26 millions de paramètres sur une boîte d’eau de 192 atomes et une supercellule de 17 280 atomes, ainsi qu’un grand modèle de 255 millions de paramètres sur le système de 192 atomes. Chaque condition est comparée à sa propre exécution FP32 native. Avec BF16x9 sur Blackwell B300, les gros noyaux GEMM du système de 17 280 atomes sont 2,2 à 2,5 fois plus rapides. Le pas complet de dynamique moléculaire y gagne 1,7 fois avec la stratégie performant et 1,4 fois avec eager. Sur 192 atomes, les gains du pas complet ne sont que de 1,2 et 1,1 fois. Le chiffre de 2,5 fois pour un noyau n’est pas le gain de la simulation entière.

Sur Ampere A100, TF32x3 utilise un noyau CUTLASS `OpMultiplyAddFastF32` et un module `TF32Linear` qui remplace des couches linéaires compatibles. Les entrées et sorties restent en FP32. Les GEMM majeurs du grand système gagnent environ 1,7 à 1,9 fois face à FP32, tandis que le pas complet gagne environ 1,5 fois. Avec le modèle de 255 millions de paramètres sur 192 atomes, le pas gagne 1,4 fois contre 1,2 fois pour le petit modèle de cette taille. NVIDIA indique qu’une seule passe TF32 peut être plus rapide, mais dégrade davantage la précision et la stabilité des trajectoires. La voie arithmétique la plus rapide n’est donc pas toujours la voie de simulation utile.

NVIDIA examine aussi la cohérence physique. Une expérience éloigne une molécule de CO d’une surface Pt(111) et observe la courbe d’énergie; une autre exécute 30 picosecondes de simulation NVE pour une boîte d’eau de 192 atomes. TF32x3 suit une courbe lisse proche de FP32 et des fluctuations d’énergie bornées comparables, tandis que TF32 en une passe montre des variations ponctuelles ou une dérive systématique plus marquée. Ce sont des résultats sur des systèmes choisis, pas une garantie de conservation pour tous les matériaux ou toutes les durées.

La relaxation d’une plaque d’alliage de 539 atomes fournit un autre test. Les trois répétitions FP32 convergent en 218 étapes; TF32x3 atteint le seuil de force de 1,0 × 10⁻³ eV/Å en 218 à 222 étapes. Les trois essais TF32 en une passe ne convergent pas dans la limite de 500 étapes, et la source indique qu’ils échouent encore quand la limite est portée à 1 000. TF32x3 accélère chaque étape par 1,5 tout en gardant la convergence observée; une passe TF32 peut accélérer l’étape par 2 sans atteindre le critère. Comparer uniquement les secondes par étape inverserait la conclusion pratique.

Avant d’adopter la méthode, il faut mesurer la part des grands GEMM dans sa propre charge et contrôler énergie, forces, trajectoires et convergence. Les petits systèmes profitent moins du changement si la recherche de voisins, les graphes et le lancement dominent. Le matériel et le réglage des noyaux comptent aussi. NVIDIA précise qu’elle n’a pas mesuré l’effet combiné avec d’autres optimisations. Notre rédaction n’a pas reproduit les essais PET. Chaque ratio rapporté reste lié au fournisseur, au GPU, au modèle, au nombre d’atomes, à la métrique et à sa référence FP32 native.