Accelerating finite-element-based projector augmented-wave density functional theory calculations with scalable GPU-centric computational methods

Dieser Beitrag stellt eine skalierbare, GPU-zentrierte Finite-Elemente-Projektor-erweiterte-Wellen-Methode (PAW-FE) vor, die algorithmische Innovationen wie gemischte Genauigkeitsarithmetik und Chebyshev-gefilterte Unterraumiteration nutzt, um erhebliche Beschleunigungen und exascale-fähige Leistung für großskalige, chemisch genaue Dichtefunktionaltheorie-Simulationen zu erreichen.

Das Wesentliche

Stellen Sie sich vor, Sie versuchen vorherzusagen, wie sich eine komplexe Maschine, etwa ein Automotor oder eine neue Batterietechnologie, verhalten wird. Um dies präzise zu tun, müssen Sie das Verhalten jedes einzelnen Elektrons in den Materialien verstehen, aus denen die Maschine besteht. Dies ist die Aufgabe eines Fachgebiets namens Dichtefunktionaltheorie (DFT). Es ist, als würde man versuchen, einen riesigen, kunstvollen Tanzboden zu simulieren, auf dem Milliarden von Elektronen synchron bewegt werden.

Lange Zeit hatten Wissenschaftler ein Problem: Die Simulation dieser Tänze für kleine Atomgruppen ist einfach, doch sobald man versucht, ein großes, komplexes System zu simulieren (wie ein winziges Metall-Nanopartikel oder eine verdrehte Materialschicht), gerät der Computer an seine Grenzen. Es ist, als würde man versuchen, einen Tanz für 100.000 Personen mit einer Methode zu leiten, die für 100 Personen entwickelt wurde; die Anweisungen verwickeln sich, der Speicher füllt sich, und die Simulation dauert ewig, bis sie abgeschlossen ist.

Diese Arbeit stellt eine neue, superschnelle Methode zur Durchführung solcher Simulationen vor, die speziell für moderne, leistungsstarke Computer entwickelt wurde, die GPUs verwenden (die gleichen Chips, die High-End-Videospiele und KI antreiben). Hier ist die Vorgehensweise, in einfache Konzepte zerlegt:

1. Der alte Weg vs. die neue Karte

• Der alte Weg (Ebene Wellen): Stellen Sie sich vor, Sie versuchen, eine Stadt mit einem riesigen, einheitlichen Gitter zu kartieren, bei dem jeder Quadratzoll die gleiche Größe hat. Wenn Sie ein winziges Detail sehen wollen (wie einen einzelnen Ziegelstein an einem Gebäude), müssen Sie das gesamte Gitter unglaublich fein machen, selbst für den leeren Himmel über der Stadt. Dies verschwendet eine massive Menge an Rechenleistung. So funktioniert die meisten aktuelle Software.
• Der neue Weg (Finite Elemente): Die Autoren verwenden einen Ansatz einer „intelligenten Karte". Stellen Sie sich eine Karte vor, die nur dort heranzoomt, wo es nötig ist (wie das belebte Stadtzentrum), und dort herauszoomt, wo es leer ist (wie der Himmel). Dies wird als Finite-Elemente-(FE)-Diskretisierung bezeichnet. Es ermöglicht ihnen, ihre Rechenleistung genau dort zu konzentrieren, wo die Elektronen interessante Dinge tun, und spart enorme Mengen an Zeit und Speicher.

2. Der „PAW"-Trick: Das magische Kostüm

Um die Mathematik noch einfacher zu machen, verwenden sie eine Methode namens Projector Augmented-Wave (PAW).

• Das Problem: Elektronen nahe dem Zentrum eines Atoms (dem Kern) wackeln und vibrieren wild, was ihre Berechnung schwierig macht.
• Die Lösung: PAW ist, als würde man den Elektronen ein „glattes Kostüm" anziehen. Es tut so, als wären die Elektronen für den Großteil der Berechnung glatt und einfach zu handhaben, behält aber einen geheimen „magischen Trick" bei, der es ihm ermöglicht, das wahre, wilde Verhalten der Elektronen genau dann sofort offenzulegen, wenn es die Details nahe dem Kern überprüfen muss. Dies erlaubt ihnen, eine viel gröbere (einfachere) Karte zu verwenden, ohne an Genauigkeit zu verlieren.

3. Der GPU-Geschwindigkeitsschub: Das Fließband

Die Autoren haben nicht nur die Karte geändert; sie haben die Art und Weise, wie der Computer die Mathematik durchführt, an moderne GPUs angepasst.

• Der Flaschenhals: Normalerweise verbringen Computer viel Zeit damit zu warten, dass Daten zwischen Speicher und Prozessor bewegt werden.
• Die Lösung: Sie haben die Mathematik so neu gestaltet, dass der Computer viele Berechnungen gleichzeitig durchführen kann (wie ein Fließband) und nicht nacheinander. Sie verwendeten zudem eine clevere Technik namens Chebyshev-Filterung, die wie ein Sieb funktioniert, das schnell die „wichtigen" Elektronen von den „unwichtigen" trennt, damit der Computer keine Zeit mit denjenigen verschwendet, die er nicht benötigt.

4. Die „gut genug"-Abkürzungen (Gemischte Genauigkeit)

Dies ist vielleicht der kreativste Teil.

• Die Analogie: Stellen Sie sich vor, Sie malen ein riesiges Wandgemälde. Für den Hintergrundhimmel müssen Sie die Farbe nicht mit mikroskopischer Präzision mischen; eine „gut genug" Mischung funktioniert gut und ist viel schneller. Sie benötigen nur extreme Präzision für die winzigen Details eines Gesichts.
• Die Anwendung: Die Autoren erkannten, dass sie für die Teile der Berechnung, bei denen es nur darum geht, die allgemeine Form richtig zu bekommen, Mathematik mit niedrigerer Genauigkeit verwenden können (wie ein Lineal mit weniger Markierungen). Dies ist auf modernen Chips viel schneller. Sie wechseln nur für die letzten, kritischen Schritte zur „ultrapräzisen" Mathematik.
• Das Ergebnis: Durch die Mischung von Mathematik mit hoher und niedriger Genauigkeit und durch das Überlappen von Datentransfers mit Berechnungen (zwei Dinge gleichzeitig tun) ließen sie die Simulation 8 bis 20 Mal schneller laufen als zuvor.

5. Was sie tatsächlich erreicht haben

Die Arbeit behauptet, dass mit diesen neuen Methoden:

• Geschwindigkeit: Sie nun Systeme mit 10.000 bis 130.000 Elektronen in einer praktischen Zeitspanne (Minuten bis Stunden) auf Supercomputern simulieren können.
• Vergleich: Ihre Methode ist für Systeme dieser Größe etwa 8 Mal schneller als die führende Standardsoftware (Quantum ESPRESSO).
• Skala: Sie führten erfolgreich eine Simulation eines „verdrehten Bilayer"-Materials (zwei Atomlagen, die miteinander verdreht sind) durch, das 130.000 Elektronen enthält. Dies ist eine Größe, die zuvor mit dieser Genauigkeit mit Standardmethoden nicht simuliert werden konnte.

Zusammenfassung

Kurz gesagt bauten die Autoren einen neuen, hocheffizienten Motor zur Simulation von Materialien. Sie kombinierten eine „intelligente Karte", die nur dort heranzoomt, wo es nötig ist, einen „magischen Kostüm"-Trick zur Vereinfachung der Mathematik und einen „Vorwärts"-Modus, der für nicht-kritische Schritte eine geringere Genauigkeit verwendet. Das Ergebnis ist ein Werkzeug, das riesige, komplexe Materialien auf modernen Supercomputern in einem Bruchteil der Zeit modellieren kann, die es früher brauchte, und damit die Tür öffnet, neue Materialien für Batterien, Elektronik und Katalysatoren viel schneller zu entwickeln.

Technische Zusammenfassung

1. Problemstellung

Die genaue Modellierung komplexer Materialsysteme (z. B. Grenzflächen, Defekte, Nanocluster, verdrehte Heterostrukturen) erfordert Dichtefunktionaltheorie (DFT)-Simulationen mit $10^4$ bis $10^5$ Elektronen. Bestehende Implementierungen stoßen jedoch auf erhebliche Engpässe:

• Einschränkungen von Plane-Wave (PW): Weit verbreitete PW-PAW-Codes (z. B. VASP, Quantum ESPRESSO) verlassen sich auf Fast-Fourier-Transformationen (FFTs), die massive All-to-All-Kommunikation erfordern. Dies macht sie bandbreitengebunden und auf modernen Exascale-GPU-Architekturen schlecht skalierbar.
• Einschränkungen im Realraum: Bestehende Realraum-Methoden verwenden oft Normerhaltende Pseudopotentiale (ONCV), die eine große Anzahl von Basisfunktionen (hohe Freiheitsgrade) benötigen, um chemische Genauigkeit zu erreichen, insbesondere für Übergangsmetalle. Dies führt zu hohen Speicherbedarfen und Rechenkosten.
• Hardware-Mismatch: Aktuelle Eigenwertlöser sind für die hohe arithmetische Intensität und die Anforderungen an niedrige Latenz moderner GPU-Cluster oft ungeeignet und nutzen weder Mischpräzisionsfähigkeiten noch überlappen sie Kommunikation und Berechnung effektiv.

2. Methodik

Die Autoren stellen PAW-FE vor, eine auf Finite-Elemente (FE) diskretisierte Projector-Augmented-Wave (PAW)-Formulierung, die für Multi-Node-GPU-Architekturen optimiert ist.

A. Mathematische Formulierung

• PAW-GHEP: Die Methode löst das verallgemeinerte hermitesche Eigenwertproblem (GHEP): $H\tilde{\Psi} = S\tilde{\Psi}\Lambda$, wobei $H$ der Hamiltonoperator und $S$ die PAW-Überlappungsmatrix ist.
• Kollineare Spin-Formalismus: Die Gleichungen werden im Rahmen eines kollinearen Spinsystems abgeleitet, um magnetische Systeme zu behandeln.
• Randbedingungen: Die FE-Diskretisierung nimmt periodische, halb-periodische (Platten) und nicht-periodische (Nanocluster) Randbedingungen auf natürliche Weise auf, ohne künstliche Periodizitätsartefakte.

B. Rechnerische Innovationen

Um das PAW-GHEP effizient auf GPUs zu lösen, entwickelten die Autoren mehrere Schlüsselalgorithmen:

1. Residuenbasierte Chebyshev-gefilterte Unterraumiteration (R-ChFSI):

   • Anstelle der Standard-ChFSI verwenden sie eine residuenbasierte Formulierung ($R = H\tilde{\Psi} - S\tilde{\Psi}\Lambda$).
   • Hauptvorteil: Diese Formulierung ist tolerant gegenüber ungenauen Matrix-Vektor-Produkten, was die Verwendung von approximativen Inversen für die PAW-Überlappungsmatrix ($S^{-1}$) und Arithmetik mit reduzierter Präzision ermöglicht, ohne die Konvergenz zu beeinträchtigen.

2. Approximative Inverse Überlappungsmatrix:

   • Anstatt die sparse Matrix $S$ explizit zu invertieren, verwenden sie eine rechnerisch günstige Näherung ($\tilde{S}^{-1}$), die aus einer diagonalen Näherung der Massmatrix und lokalisierten Korrekturen abgeleitet ist. Dies vermeidet teure globale Inversionen.

3. Multi-Resolution-Quadratur:

   • Um atomzentrierte Integrale (Pseudo-Kern-Dichten) auf groben FE-Gittern zu handhaben, wird ein Multi-Resolution-Quadratur-Schema eingesetzt. Es verwendet verfeinerte Quadraturregeln nur innerhalb der Augmentationskugeln der Atome, während es ansonsten gröbere Regeln verwendet, um Genauigkeit ohne Gitterverfeinerung sicherzustellen.

4. Mischpräzision und Kommunikation mit niedriger Präzision:

   • Berechnung: Verwendet FP32/TF32-Arithmetik für die Chebyshev-Filterungsschritte (Matrix-Matrix-Multiplikationen) und BF16 für die Kommunikation zwischen Nachbarn.
   • Robustheit: Die Abhängigkeit des R-ChFSI-Algorithmus von Residuen stellt sicher, dass diese Präzisionsreduktionen die endgültige Genauigkeit im Doppelprecision-Format des Grundzustands nicht verschlechtern.

5. Überlappung von Berechnung und Kommunikation:

   • Die Chebyshev-Filterung wird blockweise durchgeführt. Während ein Block von Wellenfunktionen interprozessuale Kommunikation (MPI) durchläuft, führt ein anderer Block Berechnungen (GEMM-Operationen) auf der GPU durch. Dies versteckt Kommunikationslatenz, einen kritischen Faktor für die Exascale-Skalierung.

6. Dichte Operationen auf Zellebene:

   • Anstatt globale sparse Matrizen zu konstruieren, formuliert die Methode Operationen als dichte Matrix-Matrix-Multiplikationen auf Zellebene. Dies maximiert die arithmetische Intensität und die Cache-Lokalität und macht sie ideal für die GPU-Ausführung.

3. Hauptbeiträge

• PAW-FE-Formulierung: Eine neuartige, FE-diskretisierte PAW-Formulierung innerhalb eines kollinearen Spin-Formalismus, die generische Randbedingungen unterstützt.
• R-ChFSI-Eigenwertlöser: Eine Erweiterung der residuenbasierten Chebyshev-gefilterten Unterraumiteration zur Lösung des PAW-GHEP, die die Verwendung von approximativen Inversen und Mischpräzision ermöglicht.
• Skalierbare GPU-Implementierung: Eine umfassende Implementierungsstrategie mit Mischpräzisionsarithmetik (FP32/TF32/BF16), Überlappung von Berechnung und Kommunikation sowie dichter linearer Algebra auf Zellebene.
• Multi-Resolution-Integration: Ein Quadratur-Schema, das eine genaue Auswertung atomzentrierter PAW-Integrale auf groben Gittern ermöglicht.

4. Ergebnisse und Leistung

Die Methode wurde auf Supercomputern der Spitzenklasse (OLCF Frontier, ALCF Aurora, ALCF Polaris) unter Verwendung von AMD-, Intel- und NVIDIA-GPUs getestet.

• Genauigkeit: Validiert gegenüber Plane-Wave-Codes (Abinit, Quantum ESPRESSO) für Moleküle ($O_2$, $NO_2$) und Kristalle (BCC Cr). Fehler in Energie und Kräften liegen um Größenordnungen unter den Anforderungen der chemischen Genauigkeit ($O(10^{-12})$ Ha/Atom für Energie, $O(10^{-6})$ Ha/Bohr für Kräfte).
• CPU-GPU-Beschleunigung: Erreichte eine 8- bis 20-fache Beschleunigung auf GPUs im Vergleich zu reinen CPU-Ausführungen (Intel- und AMD-Architekturen).
• Vergleich mit Plane-Wave (QE):
  • Für Systeme mit ca. 10.000 Elektronen erreicht PAW-FE eine 8-fache Reduktion der minimalen Wandzeit im Vergleich zu Quantum ESPRESSO.
  • Für größere Systeme (>10.000 Elektronen) steigt die Beschleunigung weiter an, bedingt durch die Lokalität von FE-Basisfunktionen im Gegensatz zur globalen Kommunikation von PW-Methoden.
• Vergleich mit ONCV-FE:
  • PAW-FE benötigt ca. 6-mal weniger Rechenressourcen und erreicht eine ca. 5-mal schnellere Lösungszeit im Vergleich zu FE-Methoden mit normerhaltenden Pseudopotentialen (ONCV), hauptsächlich aufgrund der durch die PAW-Methode ermöglichten reduzierten Freiheitsgrade.
• Exascale-Skalierbarkeit:
  • Erfolgreiche Demonstration einer Grundzustandsberechnung für ein verdrehtes bilayer WTe2-System bestehend aus 130.000 Elektronen (11.000 Atome).
  • Erreichte eine Lösungszeit von ca. 2 Minuten pro SCF-Iteration auf 400 Knoten und belegt damit die Machbarkeit der Methode für realistische großskalige Simulationen.

5. Bedeutung

Diese Arbeit etabliert PAW-FE als eine Exascale-fähige Methode für Simulationen aus ersten Prinzipien. Durch die Überbrückung der Lücke zwischen der hohen Genauigkeit der PAW-Methode und der Effizienz von Realraum-Finite-Elementen überwindet sie die Kommunikationsengpässe von Plane-Wave-Methoden. Der erfolgreiche Einsatz von Mischpräzisions- und Überlappungsstrategien auf verschiedenen GPU-Architekturen zeigt einen Weg auf, komplexe Materialsysteme mit $10^5$ Elektronen routinemäßig zu simulieren, was neue Entdeckungen in der Katalyse, bei Batterie-Grenzflächen und in Quantenmaterialien ermöglicht, die zuvor rechnerisch nicht handhabbar waren.