FR
FlashSO(2) accélère les couches à symétrie préservée par émulation entière
Accelerating SO(2)-Equivariant Materials Models through Low-Precision Emulation · arXiv:2610.35164
En termes simples
Certains modèles de matériaux sont conçus pour que, lorsqu’on tourne une structure, leurs prédictions tournent de façon cohérente. Mais la couche qui gère cette symétrie par rotation consomme une grande partie du temps de calcul, et les calculs habituels à précision réduite peuvent briser cette cohérence.
L’article présente FlashSO(2), qui découpe les calculs matriciels en petits calculs sur des nombres entiers et regroupe des opérations voisines pour gagner du temps et de la mémoire. Dans les sorties MALOQ et UMA testées, le réglage Emu 2 a respecté la limite d’erreur de rotation indiquée. Ce résultat se limite aux sorties évaluées.
La recherche
Les modèles de matériaux conçus pour que leurs prédictions évoluent correctement quand on tourne une structure consacrent une grande part de leur calcul à une couche coûteuse. Or, les calculs à précision réduite peuvent produire des erreurs qui dépendent de l’orientation.
Cette couche peut-elle utiliser du matériel plus rapide pour les calculs entiers tout en préservant la symétrie sur laquelle ces modèles s’appuient ?
FlashSO(2) émule des produits matriciels précis sur des unités rapides de calcul entier et regroupe les opérations SO(2) voisines pour accélérer les modèles de matériaux qui préservent la symétrie.
Les auteurs ont vérifié si la couche SO(2) pouvait utiliser le matériel matriciel dédié aux nombres entiers sans perdre l’équivariance par rotation, propriété selon laquelle les sorties évoluent de façon cohérente lorsque les entrées tournent.
Ils ont mesuré le plus grand écart des sorties selon différents angles de rotation, pour les Hamiltoniens de MALOQ et les forces d’UMA. Emu 2 est resté dans la tolérance de 10⁻⁴ pour les deux cibles testées.
Ils ont ensuite comparé la couche à sa référence FP32 et mesuré des tâches au niveau des modèles. Emu 2 avec regroupement des opérations a augmenté le débit des couches dans les deux configurations testées ; l’entraînement de MALOQ et l’inférence d’UMA-M étaient également plus rapides que leurs comparaisons FP32, dans des conditions précises.
Ensemble, ces résultats montrent que l’approche peut préserver la symétrie testée tout en améliorant les performances de ces modèles.
Ces résultats proviennent de modèles MALOQ et UMA représentatifs, et la plupart des évaluations ont été menées sur un seul superpuce NVIDIA GH200.
Les gains de vitesse rapportés pour les modèles correspondent à des tailles d’entrée et à des réglages précis ; ils ne permettent donc pas d’établir que les mêmes gains seraient obtenus avec d’autres modèles, matériels ou charges de travail.
Ce que montrent les preuves
- Concilier vitesse et symétrie
FlashSO(2) est une implémentation co-conçue pour le logiciel et le matériel de la couche SO(2) utilisée dans des modèles de matériaux fondés sur des réseaux de neurones sur graphes, qui représentent les éléments et leurs relations sous forme de graphe.
Ces modèles préservent l’équivariance par rotation : leurs résultats changent de façon cohérente quand l’entrée tourne. Dans nombre de ces modèles, la couche représente une grande partie du temps d’inférence, mais les calculs natifs à précision réduite peuvent introduire des erreurs dépendant de la rotation et limiter l’usage de matériel matriciel plus rapide.
- Le fonctionnement de FlashSO(2)
Pour résoudre cette difficulté, FlashSO(2) émule les produits matriciels avec des tranches INT8 mises à l’échelle, les additionne exactement en INT32, puis reconstitue le résultat. Son réglage Emu 2 utilise deux tranches.
La mise à l’échelle par blocs attribue une échelle à des groupes d’éléments, ce qui aide à conserver la précision des petites activations et à éviter les dépassements de capacité dans les longues sommes. L’implémentation regroupe aussi les opérations autour des produits matriciels et garde les résultats intermédiaires dans des registres, au lieu de les transférer entre des noyaux distincts.
- Vérifier la contrainte de symétrie
Les auteurs ont vérifié si l’implémentation préservait l’équivariance par rotation en mesurant le plus grand écart, composante par composante, à mesure que l’angle de rotation changeait. Ils ont évalué les sorties d’Hamiltonien de MALOQ et les sorties de force d’UMA selon une tolérance de 10⁻⁴.
Parmi les formats natifs, seul FP32 respectait cette limite ; TF32 la dépassait de peu. Emu 2 restait dans la limite pour les deux cibles testées, et Emu 3 ne se distinguait pas de FP32.
- Distinguer les gains de vitesse de la couche
Une comparaison par étapes a séparé l’effet du regroupement des opérations de celui de l’émulation entière : les auteurs ont comparé la couche de référence FP32 à une version FP32 regroupée, puis cette dernière à la version Emu 2 regroupée. Face à la référence, Emu 2 avec regroupement a atteint un débit 40x supérieur pour MALOQ et 7x supérieur pour UMA-L sur GH200.
Ces résultats concernent les couches et les configurations testées, et non les performances globales des modèles.
- Vérifier les gains dans les modèles
Les auteurs ont aussi mesuré les performances au niveau des modèles. Pour l’entraînement de MALOQ avec torch.compile, Emu 2 a accéléré les étapes jusqu’à 3.93x pour la plus grande entrée qui tenait en mémoire FP32, et réduit le pic de mémoire de 58% à la taille d’entrée où FP32 n’avait plus assez de mémoire.
Pour l’inférence d’UMA-M, il a exécuté une entrée donnée de cuivre cubique à faces centrées 3.33x plus vite que FP32. Ces résultats étayent l’existence de gains au-delà de la couche isolée, dans les conditions testées.
Détails
- Auteurs de l’article original
- Denghui Lu, Alexander Maeder, Alexandros Nikolaos Ziogas, Mathieu Luisier, Manasa Kaniselvan
- Soumis à arXiv
- Licence
- arXiv non-exclusive