Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data

Cet article présente un flux de travail exascale basé sur des modèles graphiques fondamentaux pour la découverte de matériaux, capable de former des modèles multi-tâches sur plus de 544 millions de structures atomiques et de cribler 1,1 milliard de structures en 50 secondes, réduisant ainsi considérablement le temps de calcul nécessaire par rapport aux méthodes de premiers principes.

L'essentiel

Imaginez que vous êtes un architecte chargé de concevoir des millions de nouveaux matériaux pour le futur : des batteries plus puissantes, des médicaments plus efficaces ou des panneaux solaires ultra-efficaces.

Dans le monde réel, pour vérifier si un matériau fonctionne, vous devez le fabriquer en laboratoire et le tester. C'est long, coûteux et impossible de tout essayer.

Avec les méthodes de calcul traditionnelles (comme la "théorie de la fonctionnelle de la densité" ou DFT), les superordinateurs peuvent simuler ces tests sans fabriquer le matériau. C'est comme faire une simulation de crash-test pour une voiture. Mais c'est extrêmement lent. Simuler un seul matériau prendrait des heures. Pour en tester un milliard, il faudrait des siècles.

Voici où entre en jeu ce papier, qui décrit une révolution pour le prix Gordon Bell (le "Nobel" de l'informatique haute performance).

1. Le Problème : Un Chaos de Données

Les chercheurs ont voulu entraîner une intelligence artificielle (IA) pour prédire ces propriétés instantanément. Mais ils se sont heurtés à un gros problème :

• Ils avaient 16 bases de données différentes, comme 16 bibliothèques avec des livres écrits dans des langues légèrement différentes.
• Certaines données étaient très précises (comme un manuel de physique quantique), d'autres moins (comme un guide de survie).
• Certaines bibliothèques étaient immenses (des milliards de structures), d'autres minuscules.
• Si vous mélangez tout ça dans un seul grand bol, l'IA devient confuse : elle oublie les petites bibliothèques importantes et ne retient que les plus grosses.

2. La Solution : Le "Chef d'Orchestre" (HydraGNN)

L'équipe a créé un système appelé HydraGNN. Imaginez un chef d'orchestre très intelligent avec une tête centrale et 16 bras différents (comme l'hydre de la mythologie).

• La tête centrale apprend les règles fondamentales de la physique (comment les atomes s'attirent ou se repoussent). C'est la base commune à tout le monde.
• Les 16 bras sont spécialisés. Chaque bras écoute une bibliothèque spécifique. L'un est expert en chimie organique, l'autre en métaux, un autre en catalyseurs.
• L'innovation clé : Au lieu de forcer l'IA à tout apprendre d'un coup, ils ont utilisé une technique appelée "apprentissage multi-tâches". La tête centrale apprend la physique de base, et chaque bras ajuste sa compréhension pour son domaine spécifique. Cela évite que l'IA ne se noie dans les données les plus volumineuses et oublie les petites.

3. L'Entraînement : Une Course de Formule 1 sur 16 000 Moteurs

Pour entraîner ce modèle, ils ont utilisé le superordinateur Frontier (le plus puissant du monde à l'époque), avec 16 000 puces graphiques (GPU) travaillant ensemble.

• Le défi : Faire travailler 16 000 moteurs ensemble sans qu'ils ne se marchent sur les pieds.
• La solution : Ils ont créé un système de "livraison de données" ultra-rapide (ADIOS2/DDStore). Imaginez un réseau de camions de livraison qui déposent les données directement dans le garage de chaque pilote (chaque GPU) avant même que la course ne commence, pour qu'aucun moteur ne s'arrête pour attendre le carburant.
• Le résultat : Ils ont entraîné le modèle sur 544 millions de structures atomiques (des milliards de combinaisons chimiques).

4. Le Choix du Meilleur Modèle : Le Concours de Talent

Ils n'ont pas juste pris un seul modèle. Ils ont lancé un grand concours (HPO) où 6 types d'architectures différentes (EGNN, SchNet, MACE, etc.) se sont affrontés.

• Le gagnant ? Un modèle basé sur PaiNN.
• Pourquoi ? Parce qu'il était le plus rapide à apprendre tout en restant très précis. C'est comme choisir le coureur qui a la meilleure vitesse de pointe tout en ayant la meilleure endurance.

5. La Magie Finale : 1 Milliard de Tests en 50 Secondes

C'est ici que ça devient fou.

• Avant : Pour vérifier 1,1 milliard de structures atomiques avec les méthodes classiques, il faudrait 6,7 années de calcul continu sur un superordinateur.
• Après : Avec leur modèle IA entraîné, ils ont pu vérifier 1,1 milliard de structures en seulement 50 secondes.
• L'analogie : C'est la différence entre lire chaque mot d'une bibliothèque entière à la main (méthode classique) et avoir un scanner magique qui lit tous les livres en une seconde (votre IA).

6. Pourquoi c'est important pour nous ?

Ce modèle n'est pas juste un gros chiffre. Il est :

• Précis : Il fonctionne aussi bien que les méthodes lentes pour la plupart des tâches.
• Adaptable : Si vous voulez l'utiliser pour un nouveau type de matériau (par exemple, des polymères pour des implants médicaux), vous pouvez le "fine-tuner" (l'ajuster) avec très peu de nouvelles données. C'est comme prendre un expert généraliste et lui donner une spécialisation rapide.
• Portable : Il fonctionne sur n'importe quel superordinateur moderne, qu'il soit aux États-Unis ou ailleurs.

En résumé :
Cette équipe a construit un "moteur de recherche" pour la matière. Au lieu de chercher une aiguille dans une botte de foin en la touchant, ils ont créé un aimant géant qui trouve l'aiguille en une fraction de seconde. Cela ouvre la porte à la découverte de nouveaux matériaux (pour l'énergie, la santé, l'environnement) à une vitesse jamais vue auparavant, transformant des années de recherche en quelques secondes.

Résumé technique

1. Problématique et Contexte

La découverte de matériaux repose traditionnellement sur des méthodes de premiers principes, telles que la théorie de la fonctionnelle de la densité (DFT). Bien que précises, ces méthodes sont extrêmement coûteuses en calcul, ce qui les rend inapplicables au criblage à grande échelle de vastes espaces chimiques.

Les modèles d'apprentissage automatique (ML) basés sur des graphes atomistiques (GFM) et les potentiels interatomiques (MLIP) offrent une alternative rapide, mais ils se heurtent à plusieurs défis majeurs :

• Hétérogénéité et déséquilibre des données : Aucun jeu de données unique n'est fiable pour tous les systèmes (organiques, inorganiques, hybrides). Les modèles doivent être entraînés sur des corpus multi-sources et multi-fidélités (différents niveaux de précision DFT), ce qui introduit des incohérences de référence énergétique et un déséquilibre massif entre les tailles des jeux de données.
• Optimisation instable : L'entraînement conjoint sur des données déséquilibrées conduit souvent à un surajustement aux domaines dominants au détriment des données minoritaires mais scientifiquement cruciales.
• Sélection de modèle et coût : Identifier la meilleure architecture pour l'échelle exascale nécessite des campagnes d'optimisation d'hyperparamètres (HPO) massives, souvent négligées dans les travaux précédents.
• Portabilité : La difficulté de déployer des workflows robustes sur des supercalculateurs hétérogènes (Frontier, Aurora, Perlmutter) avec des architectures d'accélérateurs différentes.

2. Méthodologie et Innovations

L'équipe propose un workflow complet d'échelle exascale basé sur HydraGNN, intégrant des innovations algorithmiques, système et de mesure.

A. Architecture Multi-Tâche (Multi-Task Learning - MTL)

Au lieu d'entraîner un modèle unique sur un mélange brut de données, l'approche utilise une architecture MTL :

• Couches de passage de messages partagées : Une base commune apprend les caractéristiques d'interaction atomique fondamentales transférables.
• Têtes de sortie spécifiques par jeu de données : Chaque source de données (ou niveau de fidélité) possède sa propre tête de prédiction. Cela permet d'absorber les biais spécifiques aux données (offsets de référence, bruit d'étiquetage) sans interférer avec l'apprentissage des autres tâches.
• Pondération conditionnelle par composition : Un petit MLP apprend à pondérer dynamiquement les sorties des 16 branches lors de l'inférence en fonction de la composition chimique de la structure, permettant une interpolation fluide entre les domaines d'entraînement.

B. Pipeline de Données et Infrastructure Système

Pour gérer 544+ millions de structures sur 16 384 GPU, l'équipe a conçu un pipeline co-conçu :

• ADIOS2 et DDStore : Utilisation de ces outils pour la gestion des données et le stockage distribué, permettant de mettre en cache les fragments de données fréquemment utilisés (shards) sur le stockage local NVMe des nœuds, réduisant ainsi la latence d'E/S et la charge sur le système de fichiers partagé.
• Entraînement distribué : Utilisation de la parallélisation de données distribuée et de l'exécution par tâches (task-parallel) pour équilibrer la charge de calcul selon la taille des jeux de données, évitant les temps d'attente des nœuds.

C. Optimisation d'Hyperparamètres (HPO) à Grande Échelle

• Six campagnes HPO massives ont été lancées sur le supercalculateur Frontier (utilisant 1 024 nœuds par campagne) pour explorer six architectures de réseaux de neurones à passage de messages (MPNN) : EGNN, SchNet, DimeNet, MACE, PaiNN et PNAEq.
• L'objectif était de sélectionner non seulement le modèle le plus précis, mais aussi celui offrant le meilleur compromis entre précision, temps par époque et coût computationnel.

D. Sélection du Modèle Leader et Inférence

• Le modèle PaiNN a été identifié comme le leader, capable de converger sur 100 époques dans une fenêtre de production de 12 heures sur 2 048 nœuds.
• Optimisation de l'inférence : Une série d'optimisations a été appliquée pour le criblage : réutilisation de l'encodeur, saut de branches inutiles, gradient fusionné (calcul des forces en une seule passe), et compilation via torch.compile.
• Analyse de précision : Caractérisation systématique des compromis entre BF16, FP32 et FP64 pour l'entraînement et l'inférence.

3. Résultats Clés

Performance et Échelle

• Entraînement : Le workflow a été exécuté avec succès sur Frontier (OLCF), Aurora (ALCF) et Perlmutter (NERSC), démontrant une portabilité et une scalabilité forte/faible quasi-linéaire sur des architectures hétérogènes (AMD MI250X, Intel XPU, NVIDIA A100).
• Criblage à l'échelle du milliard : Le modèle final permet d'évaluer 1,1 milliard de structures atomistiques en 50 secondes sur 9 300 nœuds de Frontier.
  • Débit global : ~21,8 millions de structures/seconde.
  • Comparaison : Ce travail équivaut à environ 6,7 années de calcul DFT continu sur la même machine.

Précision et Adaptation (Fine-Tuning)

• Adaptation aux tâches en aval : Le modèle pré-entraîné a été fine-tuné sur 12 tâches diverses (petites molécules, cristaux inorganiques, trajectoires de phase condensée).
• Gain de performance : Pour les tâches liées à la surface d'énergie potentielle (PES), le fine-tuning du modèle pré-entraîné (avec désactivation des paramètres de la base) réduit l'erreur de validation d'un ordre de grandeur par rapport à un entraînement à partir de zéro, même avec très peu de données (ex: 150 structures).
• Précision : L'entraînement en FP64 est crucial pour la stabilité. L'inférence en FP32 introduit une erreur constante de ~0,021 eV, tandis que le BF16 atteint ~0,080 eV, ce qui reste acceptable pour de nombreux criblages mais souligne l'importance du FP64 pour les applications scientifiques rigoureuses.

Sélection de Modèle

• L'analyse HPO a révélé que PaiNN était le meilleur compromis pour l'échelle exascale, grâce à sa structure d'interaction paire qui réduit le trafic mémoire et la synchronisation par rapport à des modèles plus complexes comme DimeNet ou MACE, permettant un nombre d'époques plus élevé dans le même budget temps.

4. Signification et Impact

Ce travail représente une avancée majeure pour la science des matériaux computationnelle :

1. Démocratisation du criblage à grande échelle : Il rend possible l'exploration d'espaces chimiques vastes (milliards de candidats) qui étaient auparavant inaccessibles aux méthodes de premiers principes en raison de contraintes de temps et de coût.
2. Preuve de concept pour les modèles de fondation (Foundation Models) : Il démontre qu'un seul modèle pré-entraîné sur des données hétérogènes et déséquilibrées peut servir de base robuste pour une multitude de tâches en aval, surpassant les modèles spécialisés entraînés de zéro.
3. Workflow Exascale Reproductible : Il établit un standard pour la portabilité des workflows ML sur des supercalculateurs hétérogènes, en intégrant efficacement la gestion des données, l'optimisation d'hyperparamètres et l'entraînement distribué.
4. Outil Pratique : Le modèle résultant est compact (~12 millions de paramètres, 92 Mo), permettant un déploiement direct dans des flux de travail de simulation physique sans compression excessive.

En résumé, cette recherche transforme la découverte de matériaux d'un processus séquentiel et lent en un flux de travail de criblage haut débit, fiable et rapide, en exploitant pleinement la puissance des supercalculateurs exascale et des techniques d'apprentissage automatique avancées.