Efficient Coupled-Cluster Python Frameworks for Next-Generation GPUs: A Comparative Study of CuPy and PyTorch on the Hopper and Grace Hopper Architecture

Diese Studie stellt neue Batch-Algorithmen für effiziente Coupled-Cluster-Implementierungen auf NVIDIA Hopper- und Grace-Hopper-GPUs vor, die durch optimierte Tensor-Kontraktionen mit CuPy und PyTorch eine bis zu zehnfache Beschleunigung gegenüber früheren hybriden Ansätzen erreichen.

Das Wesentliche

Das große Problem: Der überfüllte Schreibtisch

Stellen Sie sich vor, Sie sind ein genialer Chemiker, der versuchen will, die komplexesten Reaktionen der Welt zu berechnen. Um das zu tun, brauchen Sie einen riesigen Schreibtisch (den Computer-Speicher). In der Vergangenheit hatten Sie nur einen kleinen Schreibtisch (den alten GPU-Speicher von 32 GB). Wenn Sie zu viele Daten (die Moleküle) auf einmal darauf legen wollten, platzte der Schreibtisch, und Sie mussten alles wieder in den Keller (den langsamen Hauptspeicher des Computers) räumen. Das machte die Arbeit extrem langsam.

Außerdem hatten Sie zwei verschiedene Werkzeuge, um die Daten zu sortieren: CuPy und PyTorch. Beide waren gut, aber man wusste nicht genau, welches besser funktionierte, besonders wenn man auf den allerneuesten, super-schnellen Computern arbeitete.

Die Lösung: Ein neuer, smarter Stapel-Plan

Die Forscher aus diesem Papier haben nun zwei Dinge getan, um das Problem zu lösen:

1. Der neue "C-Split"-Plan (Das intelligente Sortieren)
Früher haben sie die riesigen Datenberge einfach in gleich große Stücke geschnitten, wie man ein Brot in gleich dicke Scheiben schneidet. Das war okay, aber nicht perfekt.
Jetzt haben sie eine neue Methode erfunden, die sie "C-Split" nennen. Stellen Sie sich vor, Sie haben einen riesigen Stapel Bücher. Statt sie einfach zu teilen, schauen Sie sich genau an, wie schwer jedes Buch ist und wie viel Platz auf dem Schreibtisch noch frei ist.

• Sie schneiden die Stapel dynamisch zu: Wenn ein Teil sehr groß ist, machen Sie viele kleine Schnitte. Wenn er klein ist, lassen Sie ihn ganz.
• Das ist wie ein intelligenter Lagerhaus-Manager, der genau weiß, wie er die Waren stapelt, damit kein Regal überläuft und der Gabelstapler (die GPU) immer effizient fährt.

2. Der neue Super-Computer (Hopper & Grace Hopper)
Die Forscher haben ihre Experimente auf den neuesten Maschinen von NVIDIA durchgeführt: dem H100 und dem GH200.

• Der H100 ist wie ein extrem schneller Rennwagen.
• Der GH200 ist wie ein Super-LKW mit einem riesigen Kofferraum, der direkt mit dem Fahrer (der CPU) verbunden ist. Bei diesem neuen Modell ist der Kofferraum (der Speicher) so riesig (bis zu 96 GB), dass man fast die ganze Bibliothek auf einmal hineinpacken kann, ohne hin- und herlaufen zu müssen.

Der große Wettkampf: CuPy gegen PyTorch

Die Forscher haben nun getestet, welches Werkzeug (CuPy oder PyTorch) auf diesen neuen Maschinen besser ist. Das Ergebnis ist wie ein spannendes Rennen:

• Auf dem Rennwagen (H100): Hier ist PyTorch der klare Gewinner. Es ist etwa 20 % schneller als CuPy. Man könnte sagen, PyTorch fährt auf dieser Straße einfach glatter und verliert weniger Zeit beim Bremsen und Beschleunigen.
• Auf dem Super-LKW (GH200): Hier ist das Rennen entschieden. Beide Werkzeuge sind fast gleich schnell. Der riesige Kofferraum des LKWs macht den Unterschied so groß, dass die kleinen Unterschiede zwischen den Werkzeugen kaum noch zählen.

Das Ergebnis: Ein Turbo für die Wissenschaft

Durch diese neuen Methoden (den smarten Stapel-Plan und die neuen Maschinen) haben die Forscher einen 10-fachen Geschwindigkeitsschub erreicht.

• Früher: Eine Berechnung dauerte so lange wie ein ganzer Arbeitstag.
• Jetzt: Dasselbe Ergebnis ist in wenigen Stunden fertig.

Das ist wie der Unterschied zwischen einem Fahrrad und einem Hochgeschwindigkeitszug. Besonders für sehr große Moleküle (die "Riesen" unter den Molekülen), die früher gar nicht auf einem einzigen Computer gerechnet werden konnten, ist das ein Durchbruch.

Fazit für den Alltag

Stellen Sie sich vor, Sie wollen ein riesiges Puzzle zusammenlegen.

• Alt: Sie hatten einen kleinen Tisch und mussten die Teile ständig hin- und herschieben. Es dauerte ewig.
• Neu: Sie haben einen riesigen Tisch (GH200) und eine neue Anleitung (C-Split), die Ihnen sagt, welche Teile Sie zuerst legen müssen, damit nichts im Weg liegt. Außerdem haben Sie zwei verschiedene Helfer (CuPy und PyTorch) gefunden, die je nach Tischgröße unterschiedlich gut arbeiten.

Das Ergebnis? Wir können jetzt viel schneller und genauer berechnen, wie Medikamente wirken oder wie neue Materialien aussehen, was früher Jahre gedauert hätte.

Technische Zusammenfassung

Titel

Effiziente gekoppelte-Cluster-Python-Frameworks für Next-Generation-GPUs: Ein Vergleich von CuPy und PyTorch auf der Hopper- und Grace-Hopper-Architektur

1. Problemstellung

Die Berechnung von korrelierten elektronischen Strukturen mittels der Coupled-Cluster-Methode (insbesondere CCSD – Coupled-Cluster Singles and Doubles) ist rechenintensiv und skaliert formal mit $O(o^2v^4)$ (wobei $o$ die besetzten und $v$ die virtuellen Orbitale sind).

• Herausforderung: Der Hauptengpass liegt in der sogenannten "Particle-Particle Ladder"-Kontraktion. Bei großen Molekülen und Basissätzen übersteigt der Speicherbedarf für die Zwischenresultate (Tensoren) oft die Kapazität des Videospeicher (VRAM) einzelner GPUs.
• Limitierung früherer Arbeiten: In vorherigen Arbeiten (z. B. auf NVIDIA V100S mit 32 GB VRAM) mussten große Tensoren in kleine Batches zerlegt werden, um sie in den VRAM zu laden. Die verwendeten Batch-Algorithmen waren jedoch spezifisch auf die alte Hardware optimiert und nicht allgemein anwendbar.
• Ziel: Entwicklung neuer, effizienter Batch-Algorithmen für moderne GPU-Architekturen (NVIDIA H100 und Grace Hopper GH200), die eine fast vollständige Auslagerung der Berechnungen auf die GPU ermöglichen und dabei die Speicherkapazität optimal nutzen.

2. Methodik

Die Autoren haben ihre Python-basierte CCSD-Implementierung im Softwarepaket PyBEST weiterentwickelt. Der Fokus liegt auf der Nutzung von Cholesky-zerlegten Elektronenabstoßungsintegralen (ERI), um den Speicherbedarf zu senken, und der Entwicklung neuer Batch-Strategien.

A. Neue Batch-Algorithmen

Es wurden zwei Hauptansätze vorgestellt:

1. Asymmetrisches und dynamisches Splitting-Protokoll (C-split):
   • Dies ist eine Weiterentwicklung des früheren "X-split"-Ansatzes.
   • Statt die Achsen $a$ und $b$ (virtuelle Orbitale) symmetrisch und homogen zu teilen, wird nun dynamisch entschieden, welche Achse ($a$, $b$ oder $c$) wie oft geteilt wird.
   • Der Algorithmus betrachtet die Kontraktion in zwei Schritten ($xac, xbd \to acbd$ und $acbd, ecfd \to efab$) separat und passt die Batch-Größen flexibel an die verfügbare VRAM-Kapazität an.
   • Ein entscheidender Unterschied: Die Aufteilung erfolgt nun primär entlang der $c$-Achse (statt der $e$-Achse), was auf Tests mit der Grace Hopper-Architektur basiert.
2. Generisches Batch-Protokoll:
   • Ein universeller Ansatz, der beliebige Tensor-Kontraktionen (dichte oder Cholesky-zerlegte) auf die GPU auslagert.
   • Er nutzt numpy.einsum_path, um den optimalen Pfad für die Paar-Kontraktionen zu finden.
   • Nur der erste Schritt des Kontraktionspfads wird gebatcht (entlang der Achsen, die nicht summiert werden und im Ergebnis erscheinen), um die Zwischenresultate im VRAM zu halten.

B. Hardware und Software-Umgebung

• Hardware: Vergleich zwischen einer einzelnen NVIDIA H100 (Hopper-Architektur) und dem Grace Hopper Superchip (GH200). Der GH200 bietet 96 GB HBM3 und eine nahtlose, kohärente Verbindung zwischen der Arm-basierten Grace CPU und der Hopper GPU via NVLink-C2C, was Datenübertragungsengpässe minimiert.
• Bibliotheken: Direkter Vergleich der GPU-Bibliotheken CuPy und PyTorch. Beide bieten Python-Schnittstellen, die NumPy-ähnliche Syntax unterstützen, aber unterschiedliche Backend-Optimierungen haben.
• Speichermanagement: Nutzung von Caching-Allokatoren in CuPy und PyTorch, um die Kosten für cudaMalloc/cudaFree durch Batch-Operationen zu minimieren.

3. Wichtige Beiträge

• Entwicklung des C-split-Algorithmus: Ein neuartiger, asymmetrischer Batch-Ansatz, der die Speichereffizienz im Vergleich zum alten X-split signifikant verbessert.
• Generische GPU-Engine: Implementierung einer flexiblen Schnittstelle in PyBEST, die es erlaubt, dynamisch zwischen CuPy, PyTorch und CPU-Backends (NumPy) zu wechseln, ohne den Code ändern zu müssen.
• Benchmarking auf Next-Gen-Hardware: Erste umfassende Leistungsanalyse von CCSD-Berechnungen auf H100 und GH200 unter Verwendung von Python-Frameworks.
• Demonstration des "GPU-only"-Potenzials: Zeigen, dass durch die neuen Batch-Methoden fast die gesamte CCSD-Iteration auf der GPU durchgeführt werden kann, was den Overhead durch CPU-GPU-Datenübertragungen drastisch reduziert.

4. Ergebnisse

A. Synthetische Benchmarks (Kontraktionen)

• C-split vs. X-split: Der neue C-split-Algorithmus ist auf beiden Architekturen (H100 und GH200) und mit beiden Bibliotheken (CuPy, PyTorch) deutlich schneller als der alte X-split-Ansatz.
• CuPy vs. PyTorch:
  • Auf der H100 ist PyTorch dem CuPy um ca. 20 % überlegen, insbesondere bei großen Basissätzen ($N_{basis} \ge 1000$). PyTorch scheint PCIe-Overheads besser zu verbergen (durch Overlapping von Transfer und Berechnung).
  • Auf dem GH200 sind beide Bibliotheken in der Leistung sehr ähnlich, wobei CuPy bei kleineren Problemen leicht besser abschneidet.
• Skalierung: Mit dem C-split-Algorithmus und PyTorch auf H100 wurden Beschleunigungsfaktoren von bis zu 10 im Vergleich zur ursprünglichen CuPy/X-split-Implementierung erreicht.

B. Molekulare Berechnungen (CCSD)

Getestet wurden ein Wasser-Cluster $(H_2O)_{10}$, ein hydratisiertes Uracil-Dimer und ein Farbstoffmolekül (L0) mit verschiedenen Basissätzen (bis zu 1004 Basisfunktionen).

• Geschwindigkeit: Im Vergleich zur vorherigen GPU-CPU-Hybrid-Implementierung (auf V100S) wurde eine 10-fache Beschleunigung für die gesamte CCSD-Iteration erreicht.
• Vergleich mit Referenzen: Die neuen Ergebnisse übertreffen bestehende GPU-Codes wie TeraChem und Psi4 in vielen Szenarien, insbesondere bei größeren Systemen.
• Spezifische Beobachtungen:
  • Für Systeme mit $< 500$ Basisfunktionen liefert CuPy auf GH200 die besten Ergebnisse (Faktor 4–5 schneller als alte Hybrid-Implementierung).
  • Für sehr große Systeme (L0 mit cc-pVTZ, >1000 Basisfunktionen) war nur CuPy auf GH200 in der Lage, die Berechnung durchzuführen (PyTorch lief auf H100 fehl oder war langsamer).
  • Neuer Engpass: Bei sehr großen Basissätzen verschiebt sich der Engpass von den Tensor-Kontraktionen hin zu Datenvorbereitung und anderen Operationen auf der CPU-Seite (z. B. Symmetrie-Erweiterung, Extraktion von Amplituden), die noch nicht vollständig auf die GPU ausgelagert sind.

5. Bedeutung und Ausblick

• Paradigmenwechsel: Die Arbeit zeigt, dass Python-basierte Frameworks mit modernen GPU-Architekturen (GH200) konkurrenzfähig zu hochoptimierten C++/CUDA-Codes werden können, wenn die Speicherverwaltung intelligent gelöst wird.
• Hardware-Nutzung: Der GH200 Superchip erweist sich als ideal für diese Art von Berechnungen, da die hohe Bandbreite und die kohärente Speicherverwaltung zwischen CPU und GPU die Datenübertragungsprobleme lösen, die bei reinen GPU-Clustern auftreten.
• Zukünftige Arbeiten:
  • Optimierung der verbleibenden CPU-Engpässe (Datenvorbereitung).
  • Integration von Machine Learning, um automatisch die beste Backend-Wahl (CuPy vs. PyTorch) und den besten Batch-Algorithmus basierend auf Systemgröße und Hardware vorherzusagen.
  • Skalierung auf Multi-GPU-Systeme unter Ausnutzung der NVLink-Domänen auf GH200 für Systeme mit tausenden Basisfunktionen.

Fazit: Das Paper demonstriert erfolgreich, dass durch die Kombination neuer Batch-Algorithmen (C-split), moderner Hardware (H100/GH200) und flexibler Python-Bibliotheken (PyTorch/CuPy) die Effizienz von Coupled-Cluster-Berechnungen massiv gesteigert werden kann, was neue Möglichkeiten für die Simulation komplexer molekularer Systeme eröffnet.