SMC-AI: Scaling Monte Carlo Simulation to Four Trillion Atoms with AI Accelerators

Dit paper introduceert SMC-AI, een algoritme dat Monte Carlo-simulaties op AI-acceleratoren mogelijk maakt voor een recordaantal van 4 biljoen atomen, waarmee de schaalbaarheid en doorvoer aanzienlijk worden verbeterd ten opzichte van eerdere prestaties.

De kern

Stel je voor dat je een gigantische, driedimensionale legpuzzel hebt met vier biljoen stukjes. Dat is ongeveer 4.000 keer meer stukjes dan er mensen op aarde wonen. Je doel is om deze puzzel te laten bewegen, zodat de stukjes van plaats wisselen en een nieuw, stabiel patroon vormen. Dit is wat wetenschappers doen bij het simuleren van materialen op atoomniveau.

Deze paper, getiteld SMC-AI, vertelt het verhaal van hoe een team van onderzoekers deze enorme puzzel op een manier heeft opgelost die voorheen onmogelijk leek: door gebruik te maken van computerchips die speciaal zijn ontworpen voor kunstmatige intelligentie (AI), in plaats van de traditionele supercomputers.

Hier is de uitleg in simpele taal, met een paar creatieve vergelijkingen:

1. Het Probleem: De verkeerde gereedschapskist

Stel je voor dat je een meesterklokmaker bent. Je hebt een heel specifiek gereedschap nodig om een heel klein, complex tandwiel te maken.

• De oude aanpak (HPC): Je gebruikt een zware, krachtige hamer en een vijl. Het werkt, maar het is traag en zwaar.
• De nieuwe aanpak (AI-chips): De wereld is veranderd. Nu bouwen fabrikanten vooral "snijmachines" die ontworpen zijn om miljoenen foto's in een seconde te herkennen (zoals bij gezichtsherkenning). Deze machines zijn razendsnel, maar ze zijn niet gemaakt voor het fijne werk van een klokmaker. Ze denken in blokken en lijnen, niet in losse, willekeurige bewegingen.

De onderzoekers wilden hun "klok" (de atoomsimulatie) bouwen met deze nieuwe "snijmachines" (AI-chips zoals de Huawei Ascend NPU). Het probleem? De software die ze hadden, was geschreven voor de oude hamers en vijlen. Als je die software direct op de snijmachine probeerde te draaien, zou het een ramp zijn: het zou 83 seconden duren voor iets dat normaal 0,01 seconde duurt.

2. De Oplossing: SMC-AI (De slimme vertaler)

De onderzoekers hebben SMC-AI bedacht. Dit is geen nieuwe hamer, maar een slimme vertaler en organisator.

Stel je voor dat je een enorme zaal hebt met 4.096 robots (de AI-chips).

• De oude methode: De robots kregen een lijst met instructies: "Ga naar vakje A, pak de rode bal, ga naar vakje B, check of het veilig is, ga terug." Omdat elke robot zijn eigen weg moest vinden en constant moest wachten op de volgende instructie, ontstond er chaos en stilstand.
• De SMC-AI methode: De onderzoekers hebben de regels veranderd. In plaats van dat elke robot individueel werkt, krijgen ze nu een dubbel systeem:
  1. Ze hebben twee grote tafels (geheugen). Op tafel 1 ligt de huidige situatie, op tafel 2 ligt een kopie.
  2. Alle robots werken tegelijkertijd aan hun eigen stukje van tafel 2. Ze doen allemaal precies hetzelfde type beweging (bijvoorbeeld: "wissel bal A met bal B").
  3. Pas op het allerlaatste moment kijken ze of de beweging goed was. Als ja, vullen ze tafel 1 in met de nieuwe situatie.

Dit klinkt misschien als extra werk (want ze doen soms bewegingen die ze later weer ongedaan maken), maar het is veel sneller voor deze specifieke robots. Het zorgt ervoor dat de robots nooit hoeven te wachten en altijd in een strak ritme kunnen werken.

3. Het Resultaat: Een wereldrecord

Met deze nieuwe aanpak hebben ze iets ongelooflijks bereikt:

• Ze hebben een simulatie gedaan van 4 biljoen atomen.
• Dit is 32 keer groter dan het vorige wereldrecord.
• Ze deden dit op een cluster van 4.096 AI-chips.

Het is alsof ze eerder een stadje van 100.000 mensen konden simuleren, en nu ineens de hele aarde in één keer kunnen simuleren, terwijl ze minder energie verbruiken dan de oude methoden.

4. Waarom is dit belangrijk? (De "Computers Microscoop")

Waarom doen ze dit? Stel je voor dat je een virus wilt begrijpen of een nieuw, supersterk metaal wilt ontwerpen.

• Vroeger: Je kon alleen kijken naar heel kleine stukjes (een paar duizend atomen) of het duurde jaren om een groot stuk te simuleren.
• Nu: Met SMC-AI kunnen wetenschappers kijken naar processen die plaatsvinden op een schaal die we eerder niet konden zien. Ze kunnen zien hoe nanodeeltjes in een metaal ontstaan, wat verklaart waarom sommige materialen zo sterk zijn.

Het is alsof je van een vergrootglas bent gegaan naar een computergestuurde microscoop die het hele universum in één oogopslag kan scannen.

5. De Toekomst: Flexibiliteit

Een mooi extraatje van deze methode is dat het heel flexibel is.
Stel je voor dat je een auto bouwt. De oude methode was alsof je de motor en het chassis aan elkaar vastlijmde. Als je een nieuwe motor wilde, moest je de hele auto slopen.
Met SMC-AI hebben ze de motor (het AI-model) en het chassis (de simulatie) losgekoppeld. Je kunt nu heel makkelijk een nieuwe, slimmere "motor" (een nieuw AI-model) inbouwen zonder de hele auto te hoeven herbouwen. Dit maakt het veel makkelijker voor de toekomst om nog betere simulaties te maken.

Samenvatting

De onderzoekers hebben een manier gevonden om de kracht van AI-chips (die normaal alleen voor foto's en chatbots worden gebruikt) te gebruiken voor wetenschappelijke puzzels (materiaalwetenschap). Door slimme trucjes te gebruiken om de software aan te passen aan de hardware, hebben ze een simulatie van 4 biljoen atomen mogelijk gemaakt. Dit opent de deur naar het ontwerpen van nieuwe materialen voor de toekomst, zoals sterkere auto's, betere batterijen of nieuwe medicijnen, allemaal berekend op computers die eigenlijk voor iets anders zijn gemaakt.

Technische samenvatting

Probleemstelling

De snelle ontwikkeling van diep learning heeft de hardwareontwerpwereld gedwongen om prioriteit te geven aan AI-taken, wat leidt tot de opkomst van gespecialiseerde AI-versnellers zoals NPUs (Neural Processing Units, bijv. Huawei Ascend) en TPUs. Hoewel deze hardware extreem efficiënt is voor AI-werklasten, vormt het een fundamentele uitdaging voor High-Performance Computing (HPC) taken zoals atomaire simulaties.

• Architecturale divergentie: AI-chips zijn geoptimaliseerd voor continue, grote data-blokken en matrixbewerkingen, terwijl traditionele HPC-algoritmen vaak onregelmatige geheugentoegang, fijne parallelisme en frequente vertakkingen vereisen.
• Beperkingen van bestaande methoden: De bestaande SMC-X methode (Scalable Monte Carlo at eXtreme) is zeer efficiënt op CPU's en GPU's, maar faalt bij directe porting naar AI-specifieke chips vanwege de onregelmatige geheugentoegang en de noodzaak van conditionele vertakkingen.
• Integratieproblemen: Bestaande software koppelt ML-modellen vaak te strak aan de simulatiekern, wat het moeilijk maakt om nieuwe, complexere ML-modellen (zoals MLIPs) te integreren zonder de efficiëntie te verliezen.

Methodologie: SMC-AI

De auteurs introduceren SMC-AI, een algemeen algoritmisch raamwerk dat de SMC-X methode aanpast voor AI-versnellers (NPUs en GPU's) zonder de schaalbaarheid te verliezen.

Kerninnovaties:

1. Dubbel-rooster strategie (Double-lattice strategy): Om de onregelmatige geheugentoegang van de originele SMC-X te elimineren, gebruikt SMC-AI twee roosters ($\sigma_0$ en $\sigma_1$). Het ene rooster slaat de huidige configuratie op, het andere de proefconfiguratie na een atoomuitwisseling. Dit stelt de algoritme in staat om grote, continue geheugentoegang (contiguous memory access) te realiseren, wat cruciaal is voor de efficiëntie van NPU's.
2. Ontkoppeling van ML en Simulatie: De evaluatie van het ML-model is volledig ontkoppeld van de Monte Carlo (MC) kern. De ML-modellen worden alleen aangeroepen in een specifieke functie (cal_local_energy) om lokale energieën te berekenen. Dit creëert een abstractie-laag die het eenvoudig maakt om verschillende ML-modellen (zoals qSRO of MLPNet) te integreren zonder de MC-kern aan te passen.
3. Hardware-specifieke optimalisaties voor NPU:
   • Vectorisatie via maskers: In plaats van conditionele if-statements (die inefficiënt zijn op NPUs), worden operaties uitgevoerd met behulp van bit-masks op vectorregisters.
   • Hardware-Software Parallelisme: De parallelisme-strategie is herschreven om direct te corresponderen met de AI Cube- en Vector-kernen van de Ascend-chip, waarbij data gelijkmatig wordt verdeeld over de SIMD-lanes.
   • Verbergen van geheugenlatentie: Door data vooraf te laden (prefetching) in de Unified Buffer (UB) van de AI-kernen, wordt de latency van het HBM-geheugen gemaskeerd.
   • Ghost-lagen voor Periodieke Randvoorwaarden: Een virtuele atomaire laag wordt toegevoegd om uitwisselingen over periodieke grenzen te beheren, wat noodzakelijk is omdat de SIMD-architectuur van NPUs conditionele checks inefficiënt maakt.

Belangrijkste Bijdragen

1. Recordbrekende Schaal: De implementatie van SMC-AI op een NPU-cluster heeft een simulatie van 4 biljoen (4 trillion) atomen bereikt op 4096 NPU-dies. Dit is de grootste ML-versnelde atomaire simulatie die tot nu toe is gerapporteerd.
2. Prestatieverbetering: Dit resulteert in een 32x grotere systeemgrootte en een 1,3x hogere doorvoer vergeleken met eerdere records (zoals SMC-X op GPU's), alles met een aanzienlijk kleiner computatiebudget.
3. Flexibiliteit en Portabiliteit: Het framework werkt zowel op NPUs (Huawei Ascend) als op GPU's (NVIDIA A100/H800) met uitstekende schaalbaarheid. De ontkoppeling van het ML-model maakt het een flexibel platform voor toekomstige wetenschappelijke software.
4. Nieuw ML-model (MLPNet): De auteurs stellen een nieuw, complexer ML-model voor (MLPNet) dat nauwkeuriger is dan het bestaande qSRO-model (fout van 1,78 meV vs 2,2 meV), en demonstreren dat dit model effectief kan worden gebruikt binnen het SMC-AI raamwerk.

Resultaten

• Schaalbaarheid:
  • Sterke schaalbaarheid (Strong Scaling): Op NPUs werd een efficiëntie van 82,1% bereikt bij het schalen van 128 tot 4096 dies voor een systeem van 128 miljard atomen. Op GPU's werd een efficiëntie van 96,9% bereikt.
  • Zwakke schaalbaarheid (Weak Scaling): Beide platformen vertonen bijna ideale schaalbaarheid (99,4% efficiëntie).
• Doorvoer: De NPU-implementatie bereikte een totale doorvoer van $3,77 \times 10^{10}$ atoom-stappen per seconde. Dit is ongeveer 30% van het wereldrecord op de Summit-supercomputer, maar gebruikt slechts 7,3% van het aantal versnellers (2.048 NPUs vs 27.900 V100 GPU's).
• Energie-efficiëntie: Hoewel GPU's nog steeds een iets hogere doorvoer per watt hebben, is het verschil aanzienlijk verkleind. Een enkele NPU-chip verbruikt 276 W tegenover 380 W voor een H800 GPU.
• Communicatie-Overhead: Door de implementatie van "computation-communication overlap" (het overlappen van berekening en datatransfer) wordt de communicatie- overhead gemaskeerd. Zelfs bij kubieke systemen met honderden miljarden atomen vormt communicatie slechts een klein deel van de totale runtime (bijv. 3,8% voor 16,4 miljard atomen).
• Wetenschappelijke Validatie: De simulaties van het Fe29Co29Ni28Al7Ti7 hoog-entropie legering (HEA) systeem reproduceerden correct de vorming van L12-geordende nanopartikels in een gedesordende matrix, wat in overeenstemming is met experimentele data van atoomproeftomografie.

Betekenis en Impact

Dit paper markeert een mijlpaal in de integratie van HPC en AI-hardware. Het bewijst dat het mogelijk is om gespecialiseerde AI-chips, die oorspronkelijk niet voor HPC waren ontworpen, extreem efficiënt te gebruiken voor complexe wetenschappelijke simulaties.

• Toekomst van HPC: Het biedt een blauwdruk voor hoe wetenschappelijke software kan worden gemigreerd naar de nieuwe generatie AI-centric hardware, wat essentieel is gezien de stijgende kosten van traditionele supercomputers.
• Schaalbaarheid: Het toont aan dat atomaire simulaties nu kunnen worden uitgevoerd op schalen (nanometers tot micrometers) die eerder onbereikbaar waren, waardoor wetenschappers mesoscale fenomenen met atomaire resolutie kunnen bestuderen.
• Software-architectuur: De voorgestelde abstractie tussen het ML-model en de simulatiekern vormt een fundamentele basis voor de ontwikkeling van toekomstige, schaalbare wetenschappelijke software die snel evoluerende ML-modellen kan omarmen.