Multi-GPU MBE(3)-OSV-MP2 for Performant Large-Scale ab initio Calculations

Dit artikel presenteert een performante multi-GPU-implementatie van de MBE(3)-OSV-MP2-methode die, door het overwinnen van algoritmische complexiteiten, schaalbare en snelle *ab initio*-berekeningen mogelijk maakt voor grote macromoleculen zoals insuline.

De kern

De Supercomputer in je Broodrooster: Hoe deze wetenschappers de chemie versnellen

Stel je voor dat je een gigantische puzzel moet oplossen. Maar in plaats van 1000 stukjes, heb je er miljoenen. En elke keer als je twee stukjes aan elkaar probeert te passen, moet je eerst een ingewikkelde wiskundige berekening doen om te zien of ze wel bij elkaar passen. Dit is wat chemici doen wanneer ze proberen te begrijpen hoe enorme moleculen, zoals eiwitten in ons lichaam, werken.

Deze berekeningen zijn zo zwaar dat ze normaal gesproken dagen of weken duren op de snelste supercomputers ter wereld. Dat is te lang. Als je wilt weten of een nieuw medicijn werkt, kun je niet maanden wachten op een antwoord.

Het probleem: De "Verkeersopstopping"
In de wereld van chemische berekeningen (die we ab initio noemen) is er een groot probleem. De wiskunde wordt exponentieel moeilijker naarmate het molecuul groter wordt. Het is alsof je een auto probeert te besturen, maar elke keer als je een bocht neemt, moet je eerst een heel nieuw motorblok bouwen.

Vroeger deden we dit met gewone processors (CPU's), die heel slim zijn, maar ze kunnen maar één of twee dingen tegelijk doen. Het is alsof je een berg aardappels schilt met één mes: het gaat, maar het duurt eeuwen.

De oplossing: De GPU als een leger schil-machines
De auteurs van dit paper, Qiujiang Liang en Jun Yang, hebben een oplossing gevonden. Ze gebruiken GPU's (de grafische kaarten die je ook in gaming-computers hebt).

• De analogie: Als een CPU een slimme chef-kok is die één gerecht perfect maakt, dan is een GPU een leger van 10.000 kleine koks die allemaal tegelijk een klein stukje van het gerecht snijden.
• In plaats van één grote berekening langzaam te doen, splitsen ze het probleem op in duizenden kleine stukjes die allemaal tegelijk worden berekend.

De nieuwe methode: "MBE(3)-OSV-MP2" (Laten we het de "Slimme Bouwplaat" noemen)
Deze nieuwe methode is een slimme manier om die enorme leger van koks te organiseren. Ze gebruiken drie trucjes:

1. Alleen kijken waar het belangrijk is (Lokalisatie):
   In een groot molecuul (zoals insuline) interageren atomen die ver uit elkaar liggen nauwelijks met elkaar. De oude methoden keken naar alles. Deze nieuwe methode kijkt alleen naar de atomen die dicht bij elkaar zitten. Het is alsof je in een drukke stad alleen de wegen bekijkt waar het verkeer echt vastzit, en de lege wegen negeert.
2. De "Orbitale" selectie:
   Ze kiezen slim welke elektronen ze moeten berekenen. Ze negeren de "stille" elektronen en focussen op de actieve. Dit is als het filteren van ruis in een gesprek; je hoort alleen wat er echt gezegd wordt.
3. Directe berekening zonder opslag:
   Normaal gesproken moeten computers enorme hoeveelheden data op de harde schijf opslaan en weer ophalen (zoals een bibliotheek waar je elke keer naar een ander schap moet rennen). Deze nieuwe methode doet alles direct in het snelle geheugen van de GPU. Het is alsof je de hele bibliotheek in je hoofd hebt en niet meer hoeft te rennen.

Wat hebben ze bereikt? (De resultaten)
Ze hebben dit getest op enorme moleculen:

• Waterdruppels: Ze konden een cluster van 128 watermoleculen 40 keer sneller berekenen dan de oude methoden.
• Insuline: Ze hebben de volledige berekening gedaan voor een menselijk insuline-eiwit (784 atomen!).
  • Met de oude methode zou dit dagen duren.
  • Met hun nieuwe methode op 8 krachtige GPU's duurde het 24 minuten voor de basisberekening en 6,5 uur voor de super-precieze berekening.

Waarom is dit belangrijk?
Dit is een revolutie. Het betekent dat we in de toekomst:

• Nieuwe medicijnen veel sneller kunnen ontwerpen en testen.
• Beter kunnen begrijpen hoe ziektes werken op moleculair niveau.
• Materialen kunnen ontwikkelen voor betere batterijen of zonnepanelen, zonder dat we jaren moeten wachten op de resultaten.

Kortom:
Deze wetenschappers hebben een manier gevonden om de "rekenkracht" van moderne videokaarten slim in te zetten voor de zwaarste chemische puzzels. Ze hebben de "verkeersopstopping" opgeheven en zorgen ervoor dat wetenschappers niet meer hoeven te wachten, maar direct kunnen ontdekken. Het is alsof ze de chemie van de 20e eeuw hebben versneld naar de snelheid van de 21e eeuw.

Technische samenvatting

Probleemstelling

De nauwkeurige ab initio voorspelling van de elektronische structuur van macromoleculen is computationeel zeer intensief. Hoewel de tweede-orde Møller-Plesset perturbatietheorie (MP2) essentieel is voor het nauwkeurig beschrijven van elektroncorrelatie (inclusief covalente, ionaire en van der Waals-interacties), heeft de canonieke implementatie een hoge schaalbaarheid van $O(N^5)$, waarbij $N$ de grootte van het molecuul is.

Bestaande methoden om deze complexiteit te verminderen, zoals lokale correlatiemethoden, introduceren nieuwe uitdagingen bij implementatie op Graphics Processing Units (GPUs):

1. Orbitaalllocalisatie: Het lokaliseren van bezette en virtuele moleculaire orbitalen (MO's) is vaak een $O(N^3)$ tot $O(N^4)$ bottleneck met slechte parallelle schaalbaarheid.
2. Iteratieve oplossingen: De lokale golf functie vereist het iteratief oplossen van gekoppelde residualvergelijkingen, wat complexe datatransmissie en geheugenbeheer op GPU-architecturen vereist.
3. Kleine matrixoperaties: Lokale MP2-rekeningen vertrouwen op lineaire algebra met veel kleine matrices. Standaard GPU-bibliotheken (zoals cuBLAS en cuSolver) presteren slecht bij deze kleine, latente- en geheugengebonden operaties vanwege de hoge overhead bij het starten van kernel-functies.
4. I/O en Geheugen: Het opslaan en laden van grote tussenresultaten (zoals 3c2e-integralen) veroorzaakt enorme I/O-overhead en geheugendruk, vooral bij grote systemen.

Methodologie

De auteurs presenteren een geavanceerde multi-GPU-implementatie van de MBE(3)-OSV-MP2 methode (Third-Order Many-Body Expansion met Orbital-Specific Virtuals). De aanpak combineert algoritmische optimalisaties met specifieke GPU-architectuurkenmerken:

1. Directe Integral Generator: In plaats van het opslaan van enorme 3c2e-tensoren op schijf of in het geheugen, wordt een directe generator voor half-getransformeerde integralen ($\Gamma_i$) ontwikkeld. Deze gebruikt de hoge bandbreedte van het GPU-geheugen om integralen "on-the-fly" te genereren en direct te transformeren naar lage-rang OSV-basisvectoren, waardoor de datatransferkosten worden gereduceerd tot $O(N^2)$.
2. GPU-geoptimaliseerde Localisatie:
   • Bezette Orbitalen: Implementatie van het Jacobi-Pipek-Mezey (PM) localisatie-algoritme met paarwijze rotaties, uitgevoerd op GPU-threadblokken met een schaalbaarheid van $O(N^2)$.
   • Virtuele Orbitalen (OSV): Gebruik van een gerandomiseerde SVD (rSVD) methode om de Orbitaal-Specifieke Virtuele (OSV) basis te genereren. Dit reduceert de schaalbaarheid van $O(N^4)$ naar sub-kwadratisch ($O(N^{1.5})$) en selecteert alleen de belangrijkste OSV's op basis van een drempelwaarde.
3. Aangepaste CUDA Kernels: In plaats van te vertrouwen op generieke bibliotheken, hebben de auteurs aangepaste CUDA-kernels ontwikkeld voor:
   • Het transformeren van integralen naar de OSV-basis.
   • Het berekenen van overlap- en Fock-matrices.
   • Het oplossen van de residualvergelijkingen voor MBE-clusters.
   • Deze kernels gebruiken coalesced memory access en shared memory om de latentie te minimaliseren en de parallelle efficiëntie te maximaliseren.
4. Multi-GPU Architectuur: De implementatie gebruikt MPI voor communicatie tussen CPU-nodes en NVIDIA A800 GPUs. Binnen een node wordt MPI-3 gedeeld geheugen gebruikt om dataduplicatie te voorkomen, en Remote Memory Access (RMA) minimaliseert synchronisatie-overhead tussen nodes.

Belangrijkste Bijdragen

• Eerste Multi-GPU MBE(3)-OSV-MP2: Dit is een van de eerste implementaties die lokale correlatiemethoden volledig op meerdere GPUs schaalt, specifiek gericht op het oplossen van de complexiteiten van lokale orbitalen.
• Eliminatie van I/O Bottlenecks: Door de directe regeneratie van integralen wordt de noodzaak voor massale schijf-I/O en host-device transfers voor grote tussenmatrices volledig geëlimineerd.
• Efficiëntie bij Kleine Matrices: De ontwikkeling van specifieke kernels voor kleine matrices lost het prestatieprobleem op dat traditionele GPU-bibliotheken hebben bij lokale correlatiemethoden.
• Scalabiliteit: De methode bereikt een empirische schaalbaarheid van $O(N^{1.9})$ en behoudt een hoge parallelle efficiëntie (84%) tot 24 GPUs.

Resultaten

De prestaties zijn getest op diverse systemen, variërend van waterclusters tot het menselijke insulinepeptide (784 atomen):

• Schaalbaarheid: Voor waterclusters toont de methode een schaalbaarheid van $O(N^{1.9})$, een aanzienlijke verbetering ten opzichte van de $O(N^{2.1})$ van de vorige CPU-versie.
• Snelheidswinst:
  • Tegenover Canonieke RI-MP2: Een 40-voudige versnelling in wandtijd voor $(H_2O)_{128}$ met het cc-pVDZ-basisstel.
  • Tegenover CPU-based MBE(3)-OSV-MP2: Een 10-voudige versnelling voor $(H_2O)_{190}$ met cc-pVDZ.
  • Parallelle Efficiëntie: De berekening behoudt een parallelle efficiëntie van ongeveer 90% tot 16 GPUs en daalt slechts licht tot 84% bij gebruik van 24 GPUs.
• Grote Moleculen:
  • Voor het insulinepeptide (784 atomen) met het cc-pVDZ-basisstel (7571 basisfuncties) werd de volledige MBE(3)-OSV-MP2 energie berekend in slechts 24 minuten op 8 NVIDIA A800 GPUs.
  • Met het grotere cc-pVTZ-basisstel (17448 basisfuncties) duurde de berekening 6,4 uur.
  • De piekgeheugengebruik bleef beheersbaar (88 GB voor cc-pVDZ en 241 GB voor cc-pVTZ), wat past binnen het beschikbare host-geheugen zonder dure schijf-I/O.

Betekenis

Dit werk opent nieuwe mogelijkheden voor het uitvoeren van snelle, ab initio berekeningen op echte macromoleculen (zoals eiwitten) die voorheen buiten bereik lagen voor nauwkeurige golffunctiemethoden. Door de combinatie van schaalreducerende algoritmen (MBE en OSV) met geoptimaliseerde GPU-parallelisatie, kunnen onderzoekers nu elektroncorrelatie-effecten in biologisch relevante systemen met hoge nauwkeurigheid bestuderen. De methode biedt niet alleen een enorme tijdsbesparing, maar ook economische voordelen door een efficiënter gebruik van rekenkracht vergeleken met traditionele CPU-only clusters. De auteurs wijzen erop dat toekomstig werk zich zal richten op uitbreiding naar gekoppelde cluster-methoden, periodieke systemen en analytische gradiënten.