A Survey of Neural Network Variational Monte Carlo from a Computing Workload Characterization Perspective

Dit artikel presenteert een survey en empirische GPU-characterisatie van vier neurale netwerk-variële Monte Carlo-ansätze, waarbij wordt aangetoond dat de prestaties vaak worden beperkt door rekenintensieve elementaire bewerkingen en gegevensoverdracht, wat leidt tot aanbevelingen voor algoritme-hardware co-design.

De kern

De Digitale Alchemist: Hoe een nieuwe manier van rekenen de chemie kan veranderen (en waarom het vastloopt)

Stel je voor dat je een enorme, ingewikkelde puzzel probeert op te lossen: hoe gedragen zich duizenden elektronen rondom atomen in een molecuul? Chemici noemen dit het oplossen van de "Schrödinger-vergelijking". Het is de heilige graal van de chemie, want als je dit begrijpt, kun je nieuwe medicijnen ontwerpen, betere batterijen maken of efficiëntere brandstoffen creëren.

Vroeger deden computers dit met oude, saaie formules. Maar nu proberen wetenschappers iets nieuws: Neural Network Variational Monte Carlo (NNVMC).

Wat is NNVMC? (De slimme gokker)

In plaats van met starre formules te rekenen, gebruiken deze systemen een kunstmatige intelligentie (een neurale net) die fungeert als een "slimme gokker".

• De computer probeert een oplossing te raden.
• Dan kijkt hij of die gok goed was (door te vergelijken met de natuurwetten).
• Als het niet goed is, past hij zijn "hersenen" (de parameters) een beetje aan en probeert hij opnieuw.
• Dit doet hij miljoenen keren, waarbij hij steeds nieuwe "gokken" (elektronenconfiguraties) maakt via een proces dat lijkt op het rollen van dobbelstenen (Markov Chain Monte Carlo).

Het mooie is: deze methode is heel nauwkeurig en zou theoretisch veel sneller moeten zijn dan de oude methoden. Maar in de praktijk loopt het vast. Waarom?

Het Probleem: De "Verkeersopstopping" op de Snelweg

De auteurs van dit paper hebben gekeken naar hoe deze slimme gokkers werken op moderne supercomputers (GPU's). Ze ontdekten iets verrassends.

Stel je een fabriek voor waar auto's worden gebouwd.

• De oude manier (taak 1): De machines (de GPU) zijn ontworpen om zware, zware taken te doen, zoals het gieten van enorme metalen blokken (dit noemen ze GEMM-kernels of matrixvermenigvuldigingen). Dit is waar de kracht zit.
• De NNVMC-taak: In plaats van alleen die zware blokken te gieten, moet deze fabriek ook constant kleine, snelle handelingen doen: een schroefje vastdraaien, een sticker plakken, een doosje verplaatsen, en dan weer een schroefje.

De ontdekking: De computer staat constant vast in het vastdraaien van schroefjes en het verplaatsen van doosjes.

• De zware machines (de krachtige rekenkracht) staan urenlang stil en wachten tot de kleine handelingen klaar zijn.
• De echte bottleneck is niet het rekenen, maar het verplaatsen van data. Het is alsof je een Formule-1-auto hebt, maar je rijdt erover met een fietsband. De motor schreeuwt om snelheid, maar de banden (het geheugen) kunnen niet snel genoeg data leveren.

De Vier Kampioenen (De verschillende modellen)

Het paper vergelijkt vier verschillende "gokkers" (modellen): PauliNet, FermiNet, Psiformer en Orbformer. Ze hebben allemaal hun eigen stijl, maar ze lopen allemaal tegen dezelfde muur aan, alleen op verschillende manieren:

1. PauliNet & FermiNet: Deze modellen doen heel veel "herhaling". Ze moeten een berekening doen, en dan diezelfde berekening opnieuw en opnieuw uitvoeren om de nauwkeurigheid te checken (zoals iemand die een recept blijft herschrijven om elke fout te vinden). Dit zorgt voor een enorme stroom van kleine, snelle taken die de computer verlammen.
2. Psiformer: Deze is slimmer. Hij gebruikt een andere methode (een "Transformer", zoals in ChatGPT) en doet minder herhalingen. Hierdoor is hij sneller en gebruikt hij meer van de zware rekenkracht. Maar hij is nog steeds niet perfect.
3. Orbformer: Deze probeert alles te combineren, maar door zijn complexiteit wordt hij weer erg afhankelijk van het verplaatsen van data. Hij is als een chef-kok die te veel ingrediënten van de ene naar de andere kant van de keuken moet sjouwen, waardoor hij minder tijd heeft om te koken.

De Oplossing: Samenwerken in plaats van alleen harder werken

De auteurs concluderen dat we niet zomaar "snellere computers" moeten bouwen. We moeten de manier waarop we werken veranderen. Ze stellen drie creatieve oplossingen voor:

1. De "Buurman" (PIM - Processing In Memory):
   In plaats van dat de computer (de chef) alle ingrediënten (data) moet ophalen bij de voorraadkast (het geheugen) en terugbrengt naar het aanrecht, waarom zetten we dan niet een kleine snijmachine in de voorraadkast?

   • Vergelijking: Als je een taart moet maken, loop je niet elke keer naar de koelkast voor een ei. Je zet de eierdop in de koelkast. Zo bespaar je tijd op het heen-en-weer lopen.

2. De "Wisselstaf" (Heterogene Systemen):
   Soms moet je zware blokken gieten (rekenen), en soms moet je kleine schroefjes draaien (data verplaatsen). Een enkele machine doet beide slecht.

   • Vergelijking: Gebruik een zware graafmachine voor de grondwerkzaamheden, maar wissel af met een snelle fiets voor het bezorgen van kleine pakketjes. Laat de computer weten: "Nu is het tijd om te rekenen, gebruik de zware machines. Nu is het tijd om data te verplaatsen, gebruik de snelle wegen."

3. De "Magazijn" (CPU/SSD Offload):
   Soms is de geheugenruimte op de computer gewoon te klein voor de enorme hoeveelheid data die nodig is voor grote moleculen.

   • Vergelijking: Als je keuken te klein is, huur je een opslagruimte in de kelder. Je haalt alleen de ingrediënten op die je nu nodig hebt, en de rest blijft veilig in de kelder.

Conclusie

Dit paper is een waarschuwing en een blauwdruk. Het zegt: "We bouwen de verkeerde auto's voor dit soort werk." We bouwen supercomputers die perfect zijn voor zware rekenkracht, maar NNVMC heeft juist een auto nodig die goed is in het snel verplaatsen van veel kleine dingen.

Als we de software en de hardware beter laten samenwerken (zoals de "buurman" in de voorraadkast), kunnen we binnenkort moleculen simuleren die nu onmogelijk zijn. Dat betekent snellere medicijnen, schonere energie en een beter begrip van het universum, allemaal dankzij het oplossen van een simpel verkeersprobleem in de computer.

Technische samenvatting

Titel: Een survey van Neural Network Variational Monte Carlo (NNVMC) vanuit het perspectief van werklastkarakterisering

1. Het Probleem

Het oplossen van de elektronische Schrödingervergelijking is essentieel voor kwantumchemie en materiaalkunde. Traditionele methoden (zoals CCSD(T) of FCI) bieden hoge nauwkeurigheid maar hebben een enorme rekenkosten (vaak $O(N^7)$). Neural Network Variational Monte Carlo (NNVMC) is een veelbelovend alternatief dat variational Monte Carlo combineert met neurale netwerken om de golffunctie te parametriseren, wat vaak leidt tot een betere schaalbaarheid (bijv. $O(N^4)$).

Echter, de praktische implementatie van NNVMC op moderne GPU's wordt beperkt door:

• Hoge runtime en geheugenkosten: Zelfs met efficiënte algoritmen zijn de kosten voor grotere systemen te hoog.
• Verschil met standaard AI-workloads: In tegenstelling tot taal- of visiemodellen, wordt NNVMC gedomineerd door fysica-specifieke fasen (zoals Markov Chain Monte Carlo sampling, golffunctie-construktie en Laplacian-berekening).
• Onvoldoende prestatieprognose: Het totale aantal floating-point operaties (FLOPs) is een slechte voorspeller voor de werkelijke runtime, omdat de werklast bestaat uit een heterogene mix van kernels met verschillende rekenintensiteiten.

2. Methodologie

De auteurs voeren een werklastgerichte survey en empirische karakterisering uit van vier representatieve NNVMC-ansätze (golffunctie-architecturen):

1. FermiNet en PauliNet (geïmplementeerd in de DEEPQMC-codebase).
2. Psiformer en Orbformer (geïmplementeerd in de ONEQMC-codebase).

Profilering:

• Hardware: NVIDIA RTX A5000, A100 en H200 GPU's.
• Software: JAX, CUDA 12.4, DEEPQMC en ONEQMC.
• Metingen: Gebruik van Nsight Systems en Nsight Compute om kernel-niveau metrics te verzamelen, waaronder:
  • Rekenintensiteit (Arithmetic Intensity - AI).
  • Roofline-analyse (bepalen of de kernel geheugen- of rekgebonden is).
  • Hardware-gebruik (SM-utilisatie, Tensor Core-activiteit, L2-cache hit-rate, memory throughput).
• Workload: Training en inferentie voor verschillende moleculen (van LiH tot C4H4) met een batchgrootte van 1024 elektronconfiguraties.

3. Belangrijkste Bijdragen

1. Werklastgerichte Review: Een gestructureerde vergelijking van de architecturale verschillen en uitvoeringspipelines van de vier ansätze, met name hoe ze omgaan met de Laplacian-berekening (een kritieke stap voor lokale energie).
2. Operateur- en Kernel-niveau Karakterisering: Het onthullen dat gefuseerde element-wise en data-movement kernels vaak de runtime domineren, zelfs in aanwezigheid van zware matrixvermenigvuldigingen (GEMM).
3. Hardware-gebaseerde Inzichten: Een diepgaande analyse van hoe de hardware wordt gebruikt (bijv. lage SM-utilisatie door kleine kernels) en de implicaties voor co-design van algoritme en hardware.

4. Belangrijkste Resultaten

• Runtime en Geheugen:

  • De runtime-schaalbaarheid is sterk afhankelijk van het gekozen ansatz. FermiNet en PauliNet schalen steiler dan Psiformer en Orbformer bij grotere moleculen.
  • PauliNet/FermiNet: Gebruiken een "JVP-replay" (Jacobian-Vector Product) voor de Laplacian-berekening. Dit veroorzaakt herhaalde uitvoering van eerdere fasen (Stages A-D), wat leidt tot een overvloed aan fijngranulaire, geheugengebonden kernels.
  • Psiformer/Orbformer: Gebruiken een "Hutchinson-style" schatter voor de Laplacian, wat de replay-overhead verlaagt. Hierdoor verschuift de focus meer naar sampling en grotere GEMM-kernels.

• Kernel-samenstelling en Bottlenecks:

  • PauliNet: Wordt gedomineerd door gefuseerde element-wise kernels (52% van de runtime) en layout-transformaties, vooral door de Stage E replay. De werklast is sterk geheugengebonden (lage rekenintensiteit).
  • FermiNet: Verplaatst meer runtime naar GEMM-kernels (30%) vergeleken met PauliNet, maar blijft gevoelig voor geheugenbandbreedte door de herhaalde kleine kernels.
  • Psiformer: Heeft een hogere concentratie van GEMM-kernels (62% tijdens sampling) dankzij Transformer-attention, maar de totale werklast blijft heterogeen.
  • Orbformer: Door het gebruik van FlashAttention en extra MPNN-modules (Message Passing Neural Networks) daalt het GEMM-aandeel en stijgt het aandeel van element-wise/data-movement kernels weer. Dit maakt Orbformer opnieuw sterk geheugengebonden.

• Hardware Utilisatie:

  • De meeste NNVMC-kernels opereren in het geheugengebonden gebied van het roofline-model (AI < 10 FLOP/Byte).
  • Er is een lage gebruikte instructie-throughput (vaak < 30%) en een matige L2-cache hit-rate, wat aangeeft dat data-movement en kernel-granulariteit de beperkende factoren zijn, niet de piek-rekenkracht.

5. Betekenis en Toekomstperspectief

De studie concludeert dat geoptimaliseerde NNVMC-systemen niet kunnen vertrouwen op puur rekenkracht-gerichte optimalisaties (zoals alleen het versnellen van GEMM of Attention). De resultaten onderstrepen de noodzaak van algoritme-hardware co-design met de volgende richtingen:

1. Processing-in-Memory (PIM): Omdat lage-intensiteit kernels (element-wise, reshape, transpose) de runtime domineren, zou het verplaatsen van deze berekeningen naar het geheugen (PIM) de data-movement kosten drastisch kunnen verlagen.
2. Collaboratieve GPU-PIM Systemen: Een dynamische verdeling waarbij GEMM-heavy fasen op de GPU blijven en geheugengebonden clusters worden uitbesteed aan near-memory engines.
3. Reconfigureerbare Versnelling: Hardware die zich kan aanpassen aan de verschillende fasen van de NNVMC-loop (bijv. schakelen tussen reken- en bandbreedte-geoptimaliseerde modi).
4. Beyond-Attention Optimalisatie: Voor modellen zoals Orbformer is het optimaliseren van alleen Attention-kernels onvoldoende; er is ook ondersteuning nodig voor data-movement en layout-kernels.
5. Geheugen Offloading: Voor zeer grote systemen die het GPU-geheugen overschrijden, is het offloaden van zeldzaam benaderde tensors naar CPU-geheugen of SSDs een noodzakelijke strategie.

Conclusie: De prestaties van NNVMC worden niet bepaald door het totale aantal FLOPs, maar door de efficiëntie van data-movement en de granulariteit van de kernels. Toekomstige versnelling vereist een fase-bewuste benadering die rekening houdt met de specifieke fysica-gedreven werklast van deze modellen.