Performance Benchmarking of Tensor Trains for accelerated Quantum-Inspired Homogenization on TPU, GPU and CPU architectures

Dit artikel benchmarkt Tensor Train-operaties op CPU's, GPU's en TPU's met behulp van JAX om een quantum-geïnspireerd SFFT-gebaseerd homogenisatie-algoritme aan te passen en te versnellen, waarmee succesvol hoogresolutie multiscale simulaties mogelijk worden gemaakt die variëren van 300 miljoen tot 70 miljard roosterpunten, wat onhaalbaar is met traditionele GPU-gebaseerde FFT-methoden.

De kern

Het Grote Probleem: Te Veel Data, Te Weinig Ruimte

Stel je voor dat je probeert te begrijpen hoe een complex materiaal (zoals een hoogwaardige metaallegering of een composiet) zich gedraagt onder spanning. Om dit te doen, gebruiken wetenschappers een "microscoop" om naar de minuscule interne structuur van het materiaal te kijken.

In het verleden gaven deze microscopen ons kleine, beheersbare plaatjes. Maar nieuwe technologie geeft ons nu ultra-hoge-resolutiebeelden die tientallen miljarden kleine pixels (genaamd voxels) bevatten.

Het probleem is dat het uitvoeren van de berekeningen op deze enorme afbeeldingen met traditionele methoden is alsof je een berg zand in een papieren zak probeert te dragen. De computer raakt het geheugen kwijt (de zak scheurt) of het duurt zo lang om de berekening te voltooien dat het resultaat nutteloos is tegen de tijd dat het arriveert.

De Oplossing: "Quantum-Geïnspireerde" Compressie

De auteurs stellen een nieuwe manier voor om met deze data om te gaan met behulp van een wiskundige truc genaamd Tensor Trains (TT).

Beschouw de data van het materiaal als een gigantische, 3D Rubik's Cube bestaande uit miljarden kleine blokjes.

• De Oude Manier (FFT): Proberen het probleem op te lossen door naar elk afzonderlijk blokje te kijken. Dit vereist een enorme opslagplaats om de data te bewaren en een supercomputer om de getallen te verwerken.
• De Nieuwe Manier (Tensor Trains): In plaats van elk afzonderlijk blokje op te slaan, besef je dat de kubus een patroon heeft. Je kunt de hele kubus beschrijven door slechts een paar "instructiehandleidingen" (genaamd cores) op te slaan die vertellen hoe de blokjes met elkaar verbonden zijn. Dit is vergelijkbaar met het comprimeren van een 4K-film naar een klein bestand zonder het beeld te verliezen.

Deze methode wordt "Quantum-Geïnspireerd" genoemd omdat het een techniek leent uit de kwantumfysica (de Quantum Fourier Transform) om de wiskunde op te lossen, ook al draaien de auteurs dit op gewone supercomputers en niet op echte kwantumcomputers.

Het Experiment: Wie is de Snelste Loper?

De auteurs wilden zien of deze nieuwe "gecomprimeerde" methode snel kon draaien op moderne computerchips. Ze testten drie verschillende soorten hardware:

1. CPU: De standaard hersenen van een computer (zoals een betrouwbaar, veelzijdig werkpaard).
2. GPU: Een chip ontworpen voor graphics en parallelle verwerking (zoals een team van 10.000 mieren die samenwerken).
3. TPU: Een gespecialiseerde chip gemaakt door Google specifiek voor AI (zoals een Formule 1-raceauto die gebouwd is voor één specifiek type circuit).

Ze bouwden een nieuwe motor (met behulp van een softwaretool genaamd JAX) om hun "gecomprimeerde" wiskunde op deze chips te draaien en tijden hoe snel ze gingen.

De Resultaten: Het Hangt Af van de Race

De paper concludeerde dat er geen enkele "winnaar" is. Het hangt af van de grootte van het probleem en het type wiskunde dat wordt uitgevoerd:

• Voor enorme, parallelle taken (De GPU wint): Wanneer de wiskunde bestaat uit het tegelijkertijd uitvoeren van miljoenen eenvoudige berekeningen (zoals het optellen van enorme lijsten), was de GPU het snelst. Deze schaalt prachtig op en kan enorme datasets aan die andere chips zouden laten crashen.
• Voor kleinere of complexere taken (De TPU wint): Voor bepaalde soorten wiskunde die moeilijker op te splitsen zijn, was de TPU verrassend efficiënt en versloeg deze vaak de CPU en soms zelfs de GPU.
• De CPU: Deze was het traagst, maar wel het meest stabiel. De CPU crashte niet wanneer de data te groot werd, terwijl de versnellers soms tekortkwamen aan geheugen.

Een Glitch in de Matrix:
De auteurs ontdekten een specifiek probleem met de TPU. Bij het uitvoeren van een specifiek type complexe wiskunde (genaamd SVD) met zeer grote, hoog-precieze getallen, raakte de TPU in de war en stopte deze met correct functioneren. Om dit op te lossen, moesten ze een iets tragere maar stabielere "back-up plan" (Polar Decomposition) gebruiken, specifiek voor de TPU.

Het Eindoordeel: De Limieten Doorbreken

Het meest opwindende deel van de paper is wat zij hebben bereikt met deze nieuwe opstelling:

Zij slaagden erin om homogenisatiesimulaties uit te voeren op datasets met 70 miljard gridpunten.

• De Haken en Grenzen: De beste traditionele methoden (gebruikmakend van standaard FFT) kunnen dit simpelweg niet aan. Ze raken het geheugen kwijt lang voordat ze deze omvang bereiken.
• De Doorbraak: Door de "gecomprimeerde" Tensor Train-methode op deze versnellers te gebruiken, waren zij in staat problemen op te lossen die voorheen onmogelijk waren.

Samenvatting

Beschouw deze paper als een testrit van een nieuwe, brandstofefficiënte motor (Tensor Trains) in drie verschillende auto's (CPU, GPU, TPU).

• Ze bewezen dat deze motor veel verder kan rijden (veel grotere data kan verwerken) dan de oude motoren.
• Ze ontdekten dat de GPU de beste auto is voor lange, rechte snelwegritten (enorme parallelle data).
• Ze ontdekten dat de TPU geweldig is voor specifieke, technische circuits, hoewel deze een paar eigenaardigheden heeft met hoog-precieze wiskunde.
• Belangrijker nog, ze lieten zien dat we met deze nieuwe motor eindelijk door "files" (enorme datasets) kunnen rijden die voorheen volledig geblokkeerd waren.

Technische samenvatting

Technische Samenvatting: Prestatiebenchmarking van Tensor Trains voor Quantum-geïnspireerdeerde Homogenisatie op TPU-, GPU- en CPU-architecturen

Probleemstelling
Recente vooruitgang in hoogresolutie CT-beeldvorming heeft geleid tot ultra-hoogresolutie microstructuurdatasets (die de tientallen miljarden voxels bereiken) die traditionele homogenisatiebenaderingen uitdagen. Hoewel state-of-the-art Fast Fourier Transform (FFT)-gebaseerde homogenisatietechnieken effectief zijn voor gematigde datasets, schalen hun geheugenvoetafdruk en computationele kosten als $O(dN^d \log N)$, wat ze inefficiënt maakt voor industriële problemen op grote schaal. Hoewel hardwareversnellers (GPU's en TPU's) computationele kracht bieden, overschrijden de extreme geheugeneisen van hoogresolutiegegevens vaak hun capaciteit. Hoewel Quantum Fourier Transforms (QFT) theoretische exponentiële versnellingen bieden, blijven ze onpraktisch vanwege het gebrek aan fouttolerante quantumhardware. Daarom is er behoefte aan "quantum-geïnspireerdeerde" klassieke algoritmen die gebruikmaken van laag-rang tensorrepresentaties om deze geheugen- en computationele knelpunten te overwinnen.

Methodologie
Het onderzoek onderzoekt de prestaties van het Superfast Fourier Transform (SFFT)-gebaseerde homogenisatiealgoritme, dat gebruikmaakt van Tensor Train (TT) en Tensor Train Operator (TTO) formaten om hoog-orde tensoren te representeren. De studie verloopt in twee fasen:

1. Fundamentele Operatie Benchmarking: De auteurs hebben fundamentele TT-algebra-operaties (optelling, vermenigvuldiging, contractie, orthogonalisatie en compressie) geïmplementeerd met behulp van het JAX-framework over drie hardwarearchitecturen: Dual Intel Xeon Gold 6240R CPU's, NVIDIA A100 GPU's en Google TPU v4-8. Twee implementatiemodi werden vergeleken: een "lijst-formaat" (cores opgeslagen als een lijst van arrays) en een "gebatched-formaat" (cores opgeslagen binnen een enkele gebatched array). De studie maakte gebruik van complex64 precisie om nauwkeurigheid te waarborgen, waarbij de TPU's buiten hun typische BF16-geoptimaliseerde regime werden bedreven. Prestaties werden geanalyseerd via executietijden en Roofline-modellen om te bepalen of ze zich in het geheugen-gebonden (memory-bound) of compute-gebonden regime bevinden.
2. Versnelde Homogenisatie Toepassing: De SFFT-gebaseerde homogenisatie workflow werd aangepast voor deze versnellers. Om de hoge overhead van Just-In-Time (JIT) compilatie in JAX aan te pakken wanneer tensor-ranks dynamisch veranderen, werd een "coarse-graining" strategie geïntroduceerd. Deze beperkt de tensor-ranks tot veelvouden van een basis-rank ($r_0 = 16$) om het aantal hercompilatie-gebeurtenissen te minimaliseren. Voor TPU-implementaties werd de standaard SVD-gebaseerde compressie vervangen door Polar-decompositie-gebaseerde compressie om numerieke stabiliteit te garanderen onder complex64 rekenkunde, waar SVD werd waargenomen niet te convergeren bij hoge discretisaties.

Belangrijkste Bijdragen

• Eerste Systematische TPU Benchmarking: Het artikel biedt de eerste rigoureuze benchmarking van fundamentele TT-operaties op TPU-hardware, inclusief een directe prestatievergelijking met GPU's en CPU's.
• Hardware-geaccelereerde TT Algebra: Het presenteert efficiënte implementaties van TT-algebra op moderne versnellers, waarbij de haalbaarheid van lijst-formaat versus gebatched-formaat opslag wordt geëvalueerd en specifieke prestatiekenmerken (bijv. geheugen-gebonden versus compute-gebonden gedrag) worden geïdentificeerd voor verschillende operaties.
• Praktische Implementatie van SFFT Homogenisatie: De auteurs hebben het SFFT-gebaseerde homogenisatiealgoritme succesvol aangepast voor GPU- en TPU-executie, wat simulaties van datasets mogelijk maakt variërend van 300 miljoen tot 70 miljard gridpunten—groottes die onhaalbaar zijn voor standaard GPU-gebaseerde FFT referentie-implementaties.
• Stabiliteitsanalyse: Het werk identificeert numerieke instabiliteiten in TPU-gebaseerde SVD-operaties onder complex64 precisie en stelt Polar-decompositie voor als een stabiel alternatief voor hoge discretisatie-regimes.

Resultaten

• Operatie Prestaties:
  • Parallelle Operaties: Voor zeer paralleliseerbare operaties (optelling, vermenigvuldiging, TT-TTO contractie) vertoonden GPU's een superieure schaalbaarheid bij hoge discretisatieniveaus, waarbij ze uiteindelijk de TPU's overtroffen. TPU's vertoonden een lage overhead bij lagere discretisaties maar waren strikt geheugen-gebonden over het geteste bereik.
  • Seriële Operaties: Voor seriële operaties (orthogonalisatie, compressie) presteerden TPU's over het algemeen beter dan GPU's over het volledige regime. Echter, SVD-gebaseerde compressie op TPU's slaagde er niet in te convergeren bij discretisaties rond $2^7$ onder complex64 precisie, wat de overstap naar Polar-decompositie noodzakelijk maakte.
  • Roofline Analyse: GPU's waren voornamelijk compute-gebonden voor complexe operaties, terwijl TPU's geheugen-gebonden bleven voor parallelle taken, maar overgingen naar compute-gebonden gedrag voor seriële taken bij grotere discretisaties.
• Homogenisatie Schaling:
  • De GPU-gebaseerde quantum-geïnspireerde solver schaalde succesvol op naar ongeveer 70 miljard gridpunten ($2^{18}$ punten per dimensie), wat de geheugenlimieten van de cuFFT-gebaseerde referentie-implementatie (beperkt tot $2^{12}$ punten) aanzienlijk overschrijdt.
  • CPU- en TPU-versies bereikten respectievelijk $2^{14}$ en $2^{10}$ punten per dimensie, beperkt door geheugencapaciteit.
  • Hoewel de absolute executietijden van de SFFT-methode nog niet volledig geoptimaliseerd zijn vergeleken met hoog-getunede cuFFT-bibliotheken, geeft het schaalgedrag aan dat de SFFT-aanpak uiteindelijk de FFT-gebaseerde methoden zal overtreffen naarmate de probleemomvang toeneemt, met name voor geometrieën met scheidbare structuren waar de TT-ranks gematigd blijven.
• Nauwkeurigheid: De methode behield een relatieve fout onder de 5% voor effectieve materiaaleigenschappen, gecontroleerd door de compressie-cutoff parameter.

Betekenis en Claims
Het artikel claimt een fundament te leggen voor hoogwaardige, grootschalige tensor-gebaseerde homogenisatie op moderne versnellers. Het demonstreert dat Tensor Train-technieken zowel de geheugen- als de computationele knelpunten kunnen overwinnen, wat de homogenisatie van massieve datasets mogelijk maakt die voorheen onhaalbaar waren op conventionele versnellers.

De auteurs benadrukken dat dit werk niet de fundamentele SFFT-algoritme wijzigt, maar zich richt op de efficiënte implementatie en acceleratie ervan. Ze positioneren de methode als een complementair hulpmiddel voor data-gedreven multiscale modellering, in staat om nauwkeurige referentie-oplossingen te genereren voor het trainen van neurale operatoren. De studie concludeert dat hoewel de aanpak momenteel beperkt is tot ongeveer laag-rang geometrieën (bijv. gepixeliseerde microstructuren van gelaagde composieten of rooster-materialen), het een levensvatbaar pad vertegenwoordigt naar schaalbare, natuurkundig-gebaseerde quantum-geïnspireerde solvers voor multiscale materiaalmodellering. De auteurs blijven bescheiden over de directe industriële toepasbaarheid voor willekeurige microstructuren, waarbij zij opmerken dat toekomstig werk vereist is om de numerieke stabiliteit op TPU's aan te pakken en om deze methoden uit te breiden naar hogere-orde tensornetwerken.