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

Diese Arbeit benchmarkt Tensor-Train-Operationen auf CPUs, GPUs und TPUs unter Verwendung von JAX, um einen quanteninspirierten, auf SFFT basierenden Homogenisierungsalgorithmus anzupassen und zu beschleunigen, wodurch erfolgreich hochauflösende multiskalige Simulationen im Bereich von 300 Millionen bis 70 Milliarden Gitterpunkten ermöglicht werden, die mit traditionellen GPU-basierten FFT-Methoden nicht durchführbar sind.

Das Wesentliche

Das große Problem: Zu viele Daten, zu wenig Platz

Stellen Sie sich vor, Sie versuchen zu verstehen, wie sich ein komplexes Material (wie eine hochtechnologische Metalllegierung oder ein Verbundwerkstoff) unter Belastung verhält. Um dies zu tun, nutzen Wissenschaftler ein „Mikroskop“, um in die winzige interne Struktur des Materials zu blicken.

In der Vergangenheit lieferten uns diese Mikroskope kleine, handhabbare Bilder. Aber neue Technologien liefern uns nun ultrahochauflösende Bilder, die aus Zehntausenden von Milliarden winziger Pixel (genannt Voxel) bestehen.

Das Problem ist, dass der Versuch, die Berechnungen für diese massiven Bilder mit traditionellen Methoden durchzuführen, so ist, als würde man versuchen, einen Berg Sand in einer Papiertüte zu tragen. Der Computer geht in den Speicher (die Tüte reißt) oder braucht so lange für die Berechnung, dass das Ergebnis nutzlos ist, wenn es ankommt.

Die Lösung: „Quanteninspirierte“ Kompression

Die Autoren schlagen einen neuen Weg vor, um mit diesen Daten umzugehen, indem sie einen mathematischen Trick namens Tensor Trains (TT) verwenden.

Betrachten Sie die Daten des Materials als einen riesigen, 3D-Rubik’s Cube, der aus Milliarden winziger Blöcke besteht.

• Der alte Weg (FFT): Das Problem zu lösen, indem man versucht, jeden einzelnen Block individuell zu betrachten. Dies erfordert ein riesiges Lagerhaus zur Speicherung der Daten und einen Supercomputer, um die Zahlen zu berechnen.
• Der neue Weg (Tensor Trains): Anstatt jeden einzelnen Block zu speichern, erkennt man, dass der Würfel ein Muster hat. Man kann das Ganze beschreiben, indem man nur wenige „Anleitungen“ (genannt Cores) speichert, die einem sagen, wie die Blöcke miteinander verbunden sind. Das ist vergleichbar mit der Komprimierung eines 4K-Films in eine winzige Datei, ohne das Bild zu verlieren.

Diese Methode wird als „quanteninspiriert“ bezeichnet, weil sie eine Technik aus der Quantenphysik (die Quanten-Fourier-Transformation) entlehnt, um die Mathematik zu lösen, obwohl die Autoren sie auf regulären Supercomputern und nicht auf echten Quantencomputern ausführen.

Das Experiment: Wer ist der schnellste Läufer?

Die Autoren wollten sehen, ob diese neue „komprimierte“ Methode auf modernen Computerchips schnell laufen kann. Sie testeten drei verschiedene Arten von Hardware:

1. CPU: Das Standardgehirn eines Computers (wie ein zuverlässiger, vielseitiger Arbeitspferd).
2. GPU: Ein Chip, der für Grafik und parallele Verarbeitung entwickelt wurde (wie ein Team von 10.000 Ameisen, die zusammenarbeiten).
3. TPU: Ein spezialisierter Chip von Google, der speziell für KI entwickelt wurde (wie ein Formel-1-Rennwagen, der für eine ganz bestimmte Rennstrecke gebaut wurde).

Sie entwickelten einen neuen Motor (unter Verwendung eines Software-Tools namens JAX), um ihre „komprimierte“ Mathematik auf diesen Chips auszuführen, und maßen die Zeit, die sie benötigten.

Die Ergebnisse: Es kommt darauf an

Die Studie ergab, dass es keinen einzelnen „Gewinner“ gibt. Es kommt auf die Größe des Problems und die Art der mathematischen Operation an:

• Für riesige, parallele Aufgaben (Die GPU gewinnt): Wenn die Mathematik Millionen einfacher Berechnungen gleichzeitig umfasst (wie das Aufsummieren riesiger Listen), war die GPU am schnellsten. Sie skaliert hervorragend und bewältigt massive Datensätze, an denen andere Chips scheitern würden.
• Für kleinere oder komplexere Aufgaben (Die TPU gewinnt): Für bestimmte Arten von Mathematik, die schwerer aufzuteilen sind, war die TPU überraschend effizient und schlug oft die CPU und manchmal auch die GPU.
• Die CPU: Sie war am langsamsten, aber am stabilsten. Sie stürzte nicht ab, wenn die Daten zu groß wurden, während die Beschleuniger (Accelerators) manchmal an Speicherplatzmangel litten.

Ein Fehler in der Matrix:
Die Autoren fanden ein spezifisches Problem mit der TPU. Wenn versucht wurde, eine bestimmte Art von komplexer Mathematik (genannt SVD) mit sehr großen, hochpräzisen Zahlen durchzuführen, wurde die TPU verwirrt und funktionierte nicht mehr korrekt. Um dies zu beheben, mussten sie einen etwas langsameren, aber stabileren „Plan B“ (die Polardekomposition) nur für die TPU verwenden.

Das abschließende Urteil: Grenzen sprengen

Der spannendste Teil der Arbeit ist das, was sie mit diesem neuen Setup erreicht haben:

Es gelang ihnen, Homogenisierungssimulationen auf Datensätzen mit 70 Milliarden Gitterpunkten durchzuführen.

• Der Haken: Die besten traditionellen Methoden (unter Verwendung der Standard-FFT) können dies schlichtweg nicht leisten. Sie laufen lange vor dieser Größe in den Speicherengpass.
• Der Durchbruch: Durch die Verwendung der „komprimierten“ Tensor-Train-Methode auf diesen Beschleunigern konnten sie Probleme lösen, die zuvor unmöglich waren.

Zusammenfassung

Betrachten Sie dieses Paper als eine Testfahrt mit einem neuen, treibstoffeffizienten Motor (Tensor Trains) in drei verschiedenen Autos (CPU, GPU, TPU).

• Sie haben bewiesen, dass dieser Motor viel weiter fahren kann (größere Datenmengen bewältigen) als die alten Motoren.
• Sie fanden heraus, dass die GPU das beste Auto für lange, gerade Autobahnfahrten (massive parallele Daten) ist.
• Sie fanden heraus, dass die TPU gut für spezifische, technische Strecken ist, obwohl sie einige Eigenheiten bei der Hochpräzisionsmathematik aufweist.
• Vor allem haben sie gezeigt, dass wir mit diesem neuen Motor endlich durch „Staus“ (massive Datensätze) fahren können, die früher komplett blockiert waren.

Technische Zusammenfassung

Technisches Resümee: Leistungs-Benchmarking von Tensor-Trains für quanteninspirierte Homogenisierung auf TPU-, GPU- und CPU-Architekturen

Problemstellung
Jüngste Fortschritte in der hochauflösenden CT-Bildgebung haben ultrahochaufgelöste mikrostrukturelle Datensätze generiert (die Zehnerbillarden von Voxeln erreichen), welche traditionelle Homogenisierungsansätze vor Herausforderungen stellen. Während modernste FFT-basierte (Fast Fourier Transform) Homogenisierungstechniken für moderate Datensätze effektiv sind, skalieren ihr Speicherbedarf und ihr Rechenaufwand als $O(dN^d \log N)$, was sie für industrielle Problemstellungen ineffizient macht. Obwohl Hardware-Beschleuniger (GPUs und TPUs) enorme Rechenleistung bieten, übersteigen die extremen Speicheranforderungen hochauflösender Daten oft deren Kapazität. Während Quanten-Fourier-Transformationen (QFT) einen theoretischen exponentiellen Speed-up bieten, bleiben sie aufgrund des Mangels an fehlertoleranter Quantenhardware unpraktikabel. Folglich besteht ein Bedarf an „quanteninspirierten“ klassischen Algorithmen, die niederdimensionale Tensorrepräsentationen nutzen, um diese Speicher- und Rechenengpässe zu überwinden.

Methodik
Die Arbeit untersucht die Performance des auf der Superfast Fourier Transform (SFFT) basierenden Homogenisierungsalgorithmus, welcher das Tensor-Train- (TT) und Tensor-Train-Operator-Format (TTO) nutzt, um hochdimensionale Tensoren darzustellen. Die Studie verläuft in zwei Phasen:

1. Benchmarking fundamentaler Operationen: Die Autoren implementierten fundamentale TT-Algebraoperationen (Addition, Multiplikation, Kontraktion, Orthogonalisierung und Kompression) unter Verwendung des JAX-Frameworks auf drei Hardware-Architekturen: Dual Intel Xeon Gold 6240R CPUs, NVIDIA A100 GPUs und Google TPU v4-8. Dabei wurden zwei Implementationsmodi verglichen: ein „Listen-Format“ (Kerne werden als Liste von Arrays gespeichert) und ein „Batch-Format“ (Kerne werden innerhalb eines einzelnen gebatchten Arrays gespeichert). Die Studie verwendete die komplex64-Präzision, um die Genauigkeit sicherzustellen, wobei die TPUs außerhalb ihres typischen BF16-optimierten Bereichs betrieben wurden. Die Performance wurde über Ausführungszeiten und Roofline-Modelle analysiert, um zwischen speichergebundenen (memory-bound) und rechengebundenen (compute-bound) Regimen zu unterscheiden.
2. Beschleunigte Homogenisierung als Anwendung: Der SFFT-basierte Homogenisierungs-Workflow wurde für diese Beschleuniger angepasst. Um den hohen Overhead der Just-in-Time (JIT)-Kompilierung in JAX zu adressieren, wenn sich die Tensor-Ränge dynamisch ändern, wurde eine „Vergröberungsstrategie“ (Coarse-Graining) eingeführt. Diese beschränkt die Tensor-Ränge auf Vielfache eines Basis-Rangs ($r_0 = 16$), um die Anzahl der Rekompilierungsereignisse zu minimieren. Für TPU-Implementierungen wurde die Standard-SVD-basierte Kompression durch eine auf der Polarkomposition basierende Kompression ersetzt, um die numerische Stabilität unter komplex64-Arithmetik zu gewährleisten, da die SVD bei hohen Diskretisierungsgraden beobachtbar nicht konvergierte.

Wesentliche Beiträge

• Erstes systematisches TPU-Benchmarking: Die Arbeit liefert das erste rigorose Benchmarking fundamentaler TT-Operationen auf TPU-Hardware, einschließlich eines direkten Performance-Vergleichs gegen GPUs und CPUs.
• Hardware-beschleunigte TT-Algebra: Sie präsentiert effiziente Implementierungen der TT-Algebra auf modernen Beschleunigern, evaluiert die Durchführbarkeit von Listen-Format gegenüber Batch-Format-Speicherung und identifiziert spezifische Performance-Charakteristika (z. B. speichergebundenes vs. rechengebundenes Verhalten) für verschiedene Operationen.
• Praktische Implementierung der SFFT-Homogenisierung: Die Autoren haben den SFFT-basierten Homogenisierungs-Workflow erfolgreich für die Ausführung auf GPUs und TPUs adaptiert, was die Simulation von Datensätzen im Bereich von 300 Millionen bis 70 Milliarden Gitterpunkten ermöglicht – Größen, die für Standard-GPU-basierte FFT-Referenzimplementierungen unmöglich sind.
• Stabilitätsanalyse: Die Arbeit identifiziert numerische Instabilitäten in TPU-basierten SVD-Operationen unter komplex64-Präzision und schlägt die Polarkomposition als stabile Alternative für Regime mit hoher Diskretisierung vor.

Ergebnisse

• Operations-Performance:
  • Parallele Operationen: Für hochgradig parallelisierbare Operationen (Addition, Multiplikation, TT-TTO-Kontraktion) zeigten GPUs eine überlegene Skalierbarkeit bei hohen Diskretisierungsgraden und übertrafen letztlich die TPUs. TPUs wiesen bei niedrigen Diskretisierungsgraden einen geringen Overhead auf, waren jedoch über den getesteten Bereich hinweg strikt speichergebunden.
  • Serielle Operationen: Bei seriellen Operationen (Orthogonalisierung, Kompression) übertrafen TPUs die GPUs im Allgemeinen über das gesamte Regime hinweg. Jedoch versagte die SVD-basierte Kompression auf TPUs bei Diskretisierungen um $2^7$ unter komplex64-Präzision, was den Wechsel zur Polarkomposition erforderlich machte.
  • Roofline-Analyse: GPUs waren für komplexe Operationen überwiegend rechengebunden, während TPUs für parallele Aufgaben speichergebunden blieben, sich jedoch bei seriellen Aufgaben mit größeren Diskretisierungen in Richtung rechengebundenem Verhalten entwickelten.
• Skalierung der Homogenisierung:
  • Der GPU-basierte quanteninspirierte Solver skalierte erfolgreich bis zu etwa 70 Milliarden Gitterpunkten ($2^{18}$ Punkte pro Dimension) und übertraf damit signifikant die Speichergrenzen der cuFFT-basierten Referenzimplementierung (begrenzt auf $2^{12}$ Punkte).
  • Die CPU- und TPU-Versionen erreichten $2^{14}$ bzw. $2^{10}$ Punkte pro Dimension, begrenzt durch die Speicherkapazität.
  • Obwohl die absolute Ausführungszeit der SFFT-Methode im Vergleich zu hochoptimierten cuFFT-Bibliotheken noch nicht vollständig optimiert war, deutete das Skalierungsverhalten darauf hin, dass der SFFT-Ansatz FFT-basierte Methoden bei zunehmenden Problemgrößen schließlich übertreffen würde, insbesondere für Geometrien mit separablen Strukturen, bei denen die TT-Ränge moderat bleiben.
• Genauigkeit: Die Methode hielt einen relativen Fehler von unter 5 % für effektive Materialeigenschaften ein, gesteuert durch den Kompressions-Cutoff-Parameter.

Bedeutung und Ansprüche
Die Arbeit beansprucht, ein Fundament für hochperformante, großskalige tensorbasierte Homogenisierung auf modernen Beschleunigern zu schaffen. Sie demonstriert, dass Tensor-Train-Techniken sowohl die Speicher- als auch die Rechenengpässe bei industriellen großskaligen Simulationen überwinden können, was die Homogenisierung massiver Datensätze ermöglicht, die auf konventionellen Beschleunigern zuvor nicht realisierbar waren.

Die Autoren betonen, dass diese Arbeit nicht den grundlegenden SFFT-Algorithmus modifiziert, sondern sich auf dessen effiziente Implementierung und Beschleunigung konzentriert. Sie positionieren die Methode als komplementäres Werkzeug für datengesteuerte Multiskalenmodellierung, das in der Lage ist, genaue Referenzlösungen für das Training neuronaler Operatoren zu generieren. Die Studie kommt zu dem Schluss, dass der Ansatz derzeit auf niedrigdimensionale Geometrien (z. B. pixelierte Mikrostrukturen aus geschichteten Verbundwerkstoffen oder Gittermaterialien) beschränkt ist, aber einen lebensfähigen Weg zu skalierbaren, physikbasierten quanteninspirierten Solvern für die Multiskalen-Materialmodellierung darstellt. Die Autoren bleiben hinsichtlich der unmittelbaren industriellen Anwendbarkeit für beliebige Mikrostrukturen bescheiden und merken an, dass zukünftige Arbeiten erforderlich sind, um die numerische Stabilität auf TPUs zu adressieren und diese Methoden auf höherdimensionale Tensornetzwerke zu erweitern.