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

Diese Studie stellt einen Exascale-Workflow vor, der auf HydraGNN-basierten Graph-Foundation-Modellen für die Materialentdeckung aufbaut und durch skalierbares Multi-Task-Lernen auf über 544 Millionen Strukturen sowie eine effiziente Hyperparameter-Optimierung auf dem Supercomputer Frontier die schnelle Screening-Möglichkeit von 1,1 Milliarden atomistischen Strukturen in nur 50 Sekunden ermöglicht.

Das Wesentliche

Stellen Sie sich vor, Sie sind ein genialer Architekt, der neue Materialien erfinden möchte – vielleicht einen Super-Batteriestoff für Ihr Auto oder einen unglaublich leichten Werkstoff für Flugzeuge.

Früher war das wie das Suchen nach einer Nadel im Heuhaufen, nur dass der Heuhaufen so groß ist wie das Universum und die Nadel unsichtbar ist. Um ein neues Material zu testen, mussten Wissenschaftler auf riesigen Supercomputern komplizierte physikalische Berechnungen durchführen. Das war so, als würde man jeden einzelnen Heustrohhalms einzeln mit einem Mikroskop untersuchen. Es dauerte Jahre, um nur eine Handvoll Kandidaten zu finden.

Dieses Papier beschreibt einen revolutionären neuen Weg, der diese Suche von Jahren auf 50 Sekunden verkürzt. Hier ist die Geschichte, wie sie funktioniert, einfach erklärt:

1. Der riesige Kochtopf (Die Daten)

Stellen Sie sich vor, Sie wollen den perfekten Kochkurs für alle Arten von Essen auf der Welt geben. Normalerweise würde man nur ein Rezeptbuch nehmen (z. B. nur italienisch oder nur asiatisch). Aber hier haben die Forscher 16 verschiedene Kochbücher genommen.

• Einige Bücher enthalten Rezepte für kleine, zarte Moleküle (wie ein feines Dessert).
• Andere für riesige, harte Kristalle (wie ein schwerer Braten).
• Wieder andere für Mischungen aus beidem.

Das Problem: Diese Bücher sind unterschiedlich dick (manche haben nur 10 Seiten, andere 100 Millionen), und die Autoren haben unterschiedliche Maßeinheiten verwendet. Ein normaler Koch (ein herkömmlicher Computer) würde verwirrt sein oder nur das lernen, was im dicksten Buch steht.

2. Der "Super-Lern-Koch" (Das KI-Modell)

Die Forscher haben eine spezielle KI entwickelt, die wie ein genialer Kochlehrling ist, der nicht nur rezeptiert, sondern die Grundprinzipien des Kochens versteht.

• Der gemeinsame Kern: Der Lehrling lernt zuerst die Grundlagen: "Wie verbinden sich Zutaten? Was passiert, wenn man Hitze zufügt?" Das ist der gleiche Teil für alle Rezepte.
• Die Spezialisten-Hüte: Anstatt nur einen Koch zu haben, hat dieser Lehrling 16 verschiedene Hüte auf. Wenn er ein italienisches Rezept sieht, zieht er den italienischen Hut auf. Wenn es ein Kristall-Rezept ist, zieht er den Kristall-Hut auf. So lernt er aus allen Büchern gleichzeitig, ohne dass sich die Rezepte durcheinanderbringen.

3. Der Supermarkt-Check (Die Geschwindigkeit)

Normalerweise würde man jeden neuen Heustrohhalms (jedes neue Material) einzeln im Labor testen. Das dauert ewig.
Mit diesem neuen System haben die Forscher auf dem Supercomputer "Frontier" (einem der schnellsten Computer der Welt) 1,1 Milliarden potenzielle Materialien in 50 Sekunden durchgescannt.

• Vergleich: Wenn man diese 1,1 Milliarden Materialien mit den alten Methoden testen würde, bräuchte man 6,7 Jahre ununterbrochene Rechenzeit.
• Das Ergebnis: In 50 Sekunden haben sie herausgefunden, welche Materialien vielversprechend sind. Das ist, als würde man einen ganzen Supermarkt in einer Sekunde durchsuchen, um die eine perfekte Tomate zu finden.

4. Die Feinabstimmung (Das Lernen für spezielle Aufgaben)

Was passiert, wenn man nun ein ganz spezifisches Problem hat, für das es kaum Daten gibt? (Zum Beispiel: "Wie verhält sich dieses spezielle Metall unter extremem Druck?").

• Der alte Weg: Man müsste von vorne anfangen und tausende Beispiele sammeln.
• Der neue Weg: Man nimmt den "Super-Lern-Koch", der bereits alles über die Grundlagen weiß, und gibt ihm nur ein paar wenige Beispiele für das spezielle Problem. Er passt sich sofort an und wird zum Experten für genau diese Aufgabe. Das nennt man "Fine-Tuning" (Feinabstimmung).

5. Warum das alles so wichtig ist

Früher mussten Wissenschaftler zwischen Genauigkeit (sehr langsame, teure Berechnungen) und Geschwindigkeit (schnelle, aber ungenaue Schätzungen) wählen.
Dieses neue System bricht diesen Kompromiss auf. Es ist:

• Schnell: Wie ein Blitz.
• Genau: Es behält die wissenschaftliche Präzision bei.
• Flexibel: Es funktioniert für organische Moleküle, anorganische Kristalle und alles dazwischen.

Zusammenfassung in einer Metapher

Stellen Sie sich vor, Sie suchen nach dem perfekten Musikstück für eine Party.

• Der alte Weg: Sie spielen jedes Lied der Welt einzeln ab und fragen jeden Gast, ob es ihm gefällt. Das dauert Jahre.
• Der neue Weg: Sie haben einen DJ (die KI), der die Musikgeschichte der Welt (die 1,1 Milliarden Daten) in 50 Sekunden durchgelesen hat. Er weiß genau, welche Melodien funktionieren, und kann sofort die perfekte Playlist für jede Art von Party zusammenstellen, egal ob es eine Hochzeit oder ein Rockkonzert ist.

Das Fazit: Diese Forscher haben einen Weg gefunden, wie man mit Hilfe von Supercomputern und intelligenter KI die Suche nach neuen Materialien für unsere Zukunft (bessere Batterien, sauberere Energie, leichtere Autos) von einer Ewigkeitsaufgabe in eine Sache von Sekunden verwandelt.

Technische Zusammenfassung

Titel: Exascale Multi-Task Graph Foundation Models für unausgeglichene, Multi-Fidelity-atomare Daten

1. Problemstellung

Die Entdeckung neuer Materialien, Chemikalien und molekularer Designs wird durch die hohen Kosten und den Rechenaufwand herkömmlicher First-Principles-Methoden (wie der Dichtefunktionaltheorie, DFT) behindert. Diese Methoden sind für das Screening großer chemischer Räume in industriellem Maßstab zu teuer.
Zudem bestehen in der wissenschaftlichen Machine-Learning-Forschung (ML) folgende Herausforderungen:

• Datenheterogenität: Es gibt keine einzelne Datensammlung, die für organische, anorganische und hybride Systeme gleichermaßen zuverlässig ist.
• Multi-Fidelity und Imbalance: Das Zusammenführen von Daten aus verschiedenen Quellen mit unterschiedlichen Genauigkeitsstufen (Fidelity), Approximationstheorien und Parametrisierungen führt zu inkonsistenten Energie-Referenzen, Label-Rauschen und Optimierungs-Ungleichgewichten.
• Skalierbarkeit: Bestehende Modelle können oft nicht effizient auf Exascale-Supercomputern trainiert werden, insbesondere wenn sie auf riesigen, unausgeglichenen Datensätzen basieren.
• Datenknappheit: Für viele spezifische Downstream-Aufgaben stehen nur wenige gelabelte Daten zur Verfügung, was das Training von Modellen von Grund auf erschwert.

2. Methodik und Architektur

Das Team präsentierte einen Workflow für atomare Graph-Foundation-Modelle (GFMs), die auf dem Framework HydraGNN basieren und auf dem Frontier-Supercomputer (OLCF) trainiert wurden.

• Multi-Task-Learning (MTL) Architektur:

  • Statt eines einzelnen Modells für alle Daten wird ein geteilter Message-Passing-Hintergrund (Shared Backbone) verwendet, der universelle atomare Wechselwirkungsmerkmale lernt.
  • Jeder der 16 verschiedenen Datensätze erhält einen eigenen, spezifischen Ausgabe-Head (Per-Dataset Head). Dies ermöglicht es dem Modell, datenspezifische Label-Eigenschaften und Fidelity-Offsets zu absorbieren, ohne die gemeinsamen physikalischen Merkmale zu verwässern.
  • Eine spezielle MLP-gesteuerte Branch-Weighting-Komponente gewichtet die Vorhersagen der verschiedenen Heads basierend auf der chemischen Zusammensetzung des Eingangs, was eine weiche Interpolation zwischen den Expertensystemen ermöglicht.

• Datenpipeline und Infrastruktur:

  • Nutzung von ADIOS2 und DDStore für eine skalierbare Datenbewegung. Daten werden auf NVMe-Speicher der Knoten gestaged, um Latenz zu reduzieren und das gemeinsame Dateisystem (Lustre) zu entlasten.
  • Sharded Distributed Training: HydraGNN nutzt verteilte Datenparallelität und task-parallelisierte Ausführung, um den Speicherbedarf pro Gerät zu senken und die Skalierbarkeit zu erhöhen.

• Hyperparameter-Optimierung (HPO):

  • Es wurden sechs große HPO-Kampagnen auf Frontier durchgeführt, um sechs verschiedene MPNN-Rückgrate (EGNN, SchNet, DimeNet, MACE, PaiNN, PNAEq) zu vergleichen.
  • Die Suche umfasste 574 Konfigurationen unter Berücksichtigung von Vorhersagegenauigkeit, Zeit pro Epoche und Ressourceneffizienz.

• Präzisionsanalyse:

  • Systematische Untersuchung der Auswirkungen von BF16, FP32 und FP64 auf Inferenz und Feinabstimmung (Fine-Tuning), um den Trade-off zwischen Geschwindigkeit und wissenschaftlicher Genauigkeit zu quantifizieren.

3. Wichtige Beiträge und Innovationen

• Erster Exascale-Workflow für heterogene Daten: Das erste demonstrierte Training eines atomaren GFM auf 16 offenen, hochheterogenen First-Principles-Datensätzen (insgesamt >544 Millionen Strukturen, >85 Elemente).
• Skalierbare Datenpipeline: Integration von ADIOS2/DDStore und NVMe-Staging, die das Training auf 16.384 GPUs (2.048 Knoten) unter realen Dateisystem- und Speichereinschränkungen ermöglicht.
• Architektur-Auswahl durch HPO: Statt eines festen Modells wurde die Architektur selbst als Optimierungsproblem behandelt. Das Ergebnis war ein PaiNN-basiertes Lead-Modell, das den besten Kompromiss aus Genauigkeit und Durchsatz bot.
• Composition-Conditioned Routing: Ein neuartiger Mechanismus, der es dem Modell erlaubt, dynamisch zu entscheiden, welche Datenspezialisten für eine unbekannte chemische Zusammensetzung relevant sind, ohne manuelle Zuordnung.
• Inferenz-Optimierung: Durch Techniken wie Encoder Reuse, Branch Skipping und Fused Gradients (Kombination aller 16 Vorhersagen in einem Autograd-Aufruf) sowie torch.compile wurde der Inferenz-Durchsatz massiv gesteigert.

4. Ergebnisse

• Trainingsleistung:

  • Erfolgreiches Training auf 16.384 GPUs (Frontier).
  • Das führende PaiNN-Modell (ca. 12,1 Mio. Parameter) wurde auf 2.048 Knoten für 100 Epochen trainiert.
  • Skalierbarkeit: Nahezu lineare Strong-Scaling-Effizienz bis zu 2.048 GPUs auf Perlmutter, 6.144 GPUs auf Aurora und 1.024 GPUs auf Frontier. Der Workflow ist portabel über verschiedene Supercomputer-Architekturen hinweg.

• Inferenz-Durchsatz (Screening):

  • Leistung: Evaluation von 1,1 Milliarden atomarer Strukturen in nur 50 Sekunden auf Frontier (9.300 Knoten).
  • Vergleich: Ein vergleichbares Screening mit First-Principles-Methoden würde ca. 6,7 Jahre kontinuierliche Rechenzeit benötigen.
  • Durchsatz: Ca. 21,8 Millionen Strukturen pro Sekunde (gesamtes System) bzw. 293 Strukturen/Sekunde/GPU.
  • Präzision: FP64 wurde für die Produktion gewählt, um wissenschaftliche Reproduzierbarkeit zu gewährleisten. FP32 führte zu einer konstanten Abweichung von |ΔE| ≈ 0,021 eV, BF16 zu ≈ 0,080 eV.

• Downstream-Anpassung (Fine-Tuning):

  • Das vortrainierte Modell wurde auf 12 verschiedene Downstream-Aufgaben (z. B. QM9, MD17, OQMD) feinabgestimmt.
  • Ergebnis: Für Aufgaben, die direkt mit der Potentialenergiefläche (PES) zusammenhängen, erreichte das Unfrozen Fine-Tuning (Anpassung des gesamten Modells) eine bis zu 10-fach höhere Genauigkeit als Training von Grund auf oder einfaches Fine-Tuning mit eingefrorenen Gewichten.
  • Für nicht-PES-Aufgaben (z. B. Klassifizierung) war der Vorteil geringer, was die Bedeutung von PES-Daten für das Transferlernen unterstreicht.

5. Bedeutung und Ausblick

Dieser Arbeit demonstriert, dass Exascale-Computing den Engpass der Materialentdeckung durch First-Principles-Methoden überwinden kann.

• Wissenschaftlicher Durchbruch: Sie ermöglicht das Screening chemischer Räume, die bisher unzugänglich waren, und wandelt kombinatorische Suchprobleme in machbare Workflows um.
• Praktische Anwendbarkeit: Das resultierende Modell ist kompakt (92 MB) und kann direkt in Simulations-Workflows integriert werden, ohne Kompression.
• Paradigmenwechsel: Es zeigt, dass Foundation Models für die Wissenschaft nicht nur "große Modelle" sind, sondern praktische Werkzeuge, die dateneffizientes Lernen in ressourcenarmen Szenarien ermöglichen.
• Portabilität: Der Workflow beweist, dass komplexe ML-Pipelines über heterogene Supercomputer-Systeme (Frontier, Aurora, Perlmutter) hinweg portabel und skalierbar sind, wenn Infrastruktur und Algorithmen gemeinsam entwickelt werden.

Zusammenfassend stellt diese Arbeit einen Meilenstein für die Anwendung von KI in den Materialwissenschaften dar, indem sie Skalierbarkeit, Genauigkeit und Effizienz auf Exascale-Niveau vereint.