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

Questo articolo esamina le prestazioni delle operazioni Tensor Train su CPU, GPU e TPU utilizzando JAX per adattare e accelerare un algoritmo di omogeneizzazione basato su SFFT di ispirazione quantistica, consentendo con successo simulazioni multiscala ad alta risoluzione che vanno da 300 milioni a 70 miliardi di punti griglia, altrimenti impraticabili con i tradizionali metodi FFT basati su GPU.

L'essenziale

Il Grande Problema: Troppi Dati, Poco Spazio

Immaginate di cercare di capire come si comporta un materiale complesso (come una lega metallica hi-tech o un composito) sotto stress. Per farlo, gli scienziati usano un "microscopio" per osservare la struttura interna minuscola del materiale.

In passato, questi microscopi ci fornivano immagini piccole e gestibili. Ma la nuova tecnologia ci offre oggi immagini ad altissima risoluzione che contengono decine di miliardi di minuscoli pixel (chiamati voxel).

Il problema è che cercare di eseguire i calcoli su queste immagini massicce usando i metodi tradizionali è come cercare di trasportare una montagna di sabbia in un sacchetto di carta. Il computer esaurisce la memoria (il sacchetto si rompe) o impiega così tanto tempo per calcolare che il risultato diventa inutile nel momento in cui arriva.

La Soluzione: Compressione "Ispirata alla Quantistica"

Gli autori propongono un nuovo modo per gestire questi dati utilizzando un trucco matematico chiamato Tensor Trains (TT).

Pensate ai dati del materiale come a un gigantesco Cubo di Rubik 3D composto da miliardi di piccoli blocchi.

• Il Vecchio Metodo (FFT): Cercare di risolvere il problema guardando ogni singolo blocco individualmente. Questo richiede un magazzino enorme per conservare i dati e un supercomputer per elaborare i numeri.
• Il Nuovo Metolo (Tensor Trains): Invece di memorizzare ogni singolo blocco, vi rendete conto che il cubo ha un modello. Potete descrivere l'intero oggetto memorizzando solo alcuni "manuali di istruzioni" (chiamati core) che spiegano come i blocchi si connettono tra loro. È come comprimere un film in 4K in un file minuscolo senza perdere l'immagine.

Questo metodo è chiamato "ispirato alla quantistica" perché prende in prestito una tecnica dalla fisica quantistica (la Trasformata di Fourier Quantistica) per risolvere la matematica, anche se gli autori la eseguono su normali supercomputer e non su veri computer quantistici.

L'Esperimento: Chi è il Corridore più Veloce?

Gli autori volevano vedere se questo nuovo metodo "compresso" potesse girare velocemente sui moderni chip per computer. Hanno testato tre diversi tipi di hardware:

1. CPU: Il cervello standard di un computer (come un affidabile lavoratore versatile per ogni scopo).
2. GPU: Un chip progettato per la grafica e l'elaborazione parallela (come una squadra di 10.000 formiche che lavorano insieme).
3. TPU: Un chip specializzato creato da Google specificamente per l'IA (come un'auto di Formula 1 costruita per un tipo specifico di pista).

Hanno costruito un nuovo motore (usando uno strumento software chiamato JAX) per far girare la loro matematica "compressa" su questi chip e hanno cronometrato la loro velocità.

I Risultati: Dipende dalla Gara

Il documento ha scoperto che non esiste un unico "vincitore". Dipende dalle dimensioni del problema e dal tipo di calcolo effettuato:

• Per compiti enormi e paralleli (La GPU Vince): Quando la matematica prevede l'esecuzione di milioni di calcoli semplici contemporaneamente (come sommare liste enormi), la GPU è stata la più veloce. Scalabilità in modo splendido, gestendo dataset massicci che manderebbero in crash gli altri chip.
• Per compiti più piccoli o complessi (La TPU Vince): Per certi tipi di matematica più difficili da suddividere, la TPU è stata sorprendentemente efficiente, superando spesso la CPU e talvolta la GPU.
• La CPU: È stata la più lenta, ma la più stabile. Non è andata in crash quando i dati sono diventati troppo grandi, a differenza degli acceleratori che a volte hanno esaurito la memoria.

Un Glitch nella Matrice:
Gli autori hanno riscontrato un problema specifico con la TPU. Quando cercavano di eseguire un tipo particolare di matematica complessa (chiamata SVD) su numeri molto grandi e ad alta precisione, la TPU si confondeva e smetteva di funzionare correttamente. Per risolvere il problema, hanno dovuto usare un "piano di riserva" leggermente più lento ma più stabile (la Decomposizione Polare) proprio per la TPU.

Il Verdetto Finale: Rompere i Limiti

La parte più entusiasmante del documento è ciò che hanno ottenuto con questa nuova configurazione:

Sono riusciti a eseguire simulazioni di omogeneizzazione su dataset con 70 miliardi di punti di griglia.

• Il Problema: I migliori metodi tradizionali (usando la FFT standard) semplicemente non possono farlo. Esauriscono la memoria molto prima di raggiungere tali dimensioni.
• La Svolta: Utilizzando il metodo compresso dei Tensor Trains su questi acceleratori, sono stati in grado di risolvere problemi che prima erano impossibili.

Riassunto

Pensate a questo documento come a un test drive per un nuovo motore a basso consumo (Tensor Trains) in tre auto diverse (CPU, GPU, TPU).

• Hanno dimostato che questo motore può guidare molto più lontano (gestire molti più dati) rispetto ai vecchi motori.
• Hanno scoperto che la GPU è l'auto migliore per lunghi viaggi in autostrada dritti (dati paralleli massicci).
• Hanno scoperto che la TPU è ottima per piste tecniche specifiche, nonostante abbia qualche piccola stranezza con la matematica ad alta precisione.
• Soprattutto, hanno dimostato che con questo nuovo motore possiamo finalmente attraversare i "ingorghi stradali" (dataset massicci) che prima erano completamente bloccati.

Sommario tecnico

Sintesi Tecnica: Benchmarking delle Prestazioni dei Tensor Train per l'Omogeneizzazione Quantum-Inspired su Architetture TPU, GPU e CPU

Definizione del Problema
Recenti progressi nell'imaging CT ad alta risoluzione hanno generato dataset microstrutturali ultra-elevati (che raggiungono decine di miliardi di voxel) che mettono alla prova i tradizionali approcci di omogeneizzazione. Sebbene le tecniche di omogeneizzazione basate su Fast Fourier Transform (FFT) all'avanguardia siano efficaci per dataset di medie dimensioni, il loro footprint di memoria e il costo computazionale scalano come $O(dN^d \log N)$, rendendoli inefficienti per problemi di scala industriale. Sebbene gli acceleratori hardware (GPU e TPU) offrano potenza computazionale, i requisiti estremi di memoria per i dati ad alta risoluzione spesso superano la loro capacità. Sebbene le Quantum Fourier Transforms (QFT) offrano accelerazioni esponenziali teoriche, rimangono impraticabili a causa della mancanza di hardware quantistico fault-tolerant. Di conseguenza, vi è la necessità di algoritmi classici "quantum-inspired" che sfruttino rappresentazioni tensoriali a basso rango per superare questi colli di bottiglia di memoria e computazione.

Metodologia
Il documento investiga le prestazioni dell'algoritmo di omogeneizzazione basato su Superfast Fourier Transform (SFFT), che utilizza i formati Tensor Train (TT) e Tensor Train Operator (TTO) per rappresentare tensori di alto ordine. Lo studio procede in due fasi:

1. Benchmarking delle Operazioni Fondamentali: Gli autori hanno implementato operazioni fondamentali dell'algebra TT (addizione, moltiplicazione, contrazione, ortogonalizzazione e compressione) utilizzando il framework JAX su tre architetture hardware: Dual Intel Xeon Gold 6240R CPU, NVIDIA A100 GPU e Google TPU v4-8. Sono stati confrontati due modalità di implementazione: una "list-format" (core memorizzati come una lista di array) e una "batched-format" (core memorizzati all'interno di un singolo array batched). Lo studio ha utilizzato la precisione complex64 per garantire l'accuratezza, operando le TPU al di fuori del loro tipico regime ottimizzato per BF16. Le prestazioni sono state analizzate tramite tempi di esecuzione e modelli Roofline per determinare i regimi limitati dalla memoria (memory-bound) o dal calcolo (compute-bound).
2. Applicazione dell'Omogeneizzazione Accelerata: Il workflow di omogeneizzazione basato su SFFT è stato adattato per questi acceleratori. Per affrontare l'elevato overhead della compilazione Just-In-Time (JIT) in JAX quando i ranghi dei tensori cambiano dinamicamente, è stata introdotta una strategia di "coarse-graining". Questa restringe i ranghi dei tensori a multipli di un rango base ($r_0 = 16$) per minimizzare gli eventi di ricompilazione. Per le implementazioni su TPU, la compressione standard basata su SVD è stata sostituita con la compressione basata su decomposizione Polare per garantire la stabilità numerica sotto aritmetica complex64, dove è stato osservato che la SVD non riusciva a convergere ad alte discretizzazioni.

Contributi Chiave

• Primo Benchmarking Sistematico su TPU: Il documento fornisce il primo benchmarking rigoroso delle operazioni fondamentali TT sull'hardware TPU, incluso un confronto diretto delle prestazioni rispetto a GPU e CPU.
• Algebra TT Accelerata dall'Hardware: Presenta implementazioni efficienti dell'algebra TT su moderni acceleratori, valutando la fattibilità dello storage list-format rispetto al batched-format e identificando specifiche caratteristiche di prestazione (ad esempio, comportamento memory-bound vs compute-bound) per diverse operazioni.
• Implementazione Pratica dell'Omogeneizzazione SFFT: Gli autori hanno adattato con successo l'algoritmo di omogeneizzazione basato su SFFT per l'esecuzione su GPU e TPU, consentendo la simulazione di dataset che vanno da 300 milioni a 70 miliardi di punti di griglia — dimensioni impraticabili per le implementazioni di riferimento FFT standard basate su GPU.
• Analisi della Stabilità: Il lavoro identifica instabilità numeriche nelle operazioni SVD su TPU sotto precisione complex64 e propone la decomposizione Polare come alternativa stabile per regimi di alta discretizzazione.

Risultati

• Prestazioni delle Operazioni:
  • Operazioni Parallele: Per operazioni altamente parallelizzabili (addizione, moltiplicazione, contrazione TT-TTO), le GPU hanno dimostrato una scalabilità superiore ai livelli di discretizzazione elevati, superando infine le TPU. Le TPU hanno mostrato un basso overhead a discretizzazioni inferiori, ma sono state strettamente limitate dalla memoria (memory-bound) nell'intervallo testato.
  • Operazioni Seriali: Per operazioni seriali (ortogonalizzazione, compressione), le TPU hanno generalmente superato le GPU in tutto il regime. Tuttavia, la compressione basata su SVD sulle TPU ha fallito la convergenza intorno a $2^7$ sotto precisione complex64, rendendo necessaria la transizione alla decomposizione Polare.
  • Analisi Roofline: Le GPU erano prevalentemente limitate dal calcolo (compute-bound) per le operazioni complesse, mentre le TPU rimanevano limitate dalla memoria (memory-bound) per i compiti paralleli, ma transitavano verso un comportamento limitato dal calcolo per i compiti seriali a grandi discretizzazioni.
• Scalabilità dell'Omogeneizzazione:
  • Il solver quantum-inspired basato su GPU è riuscito a scalare fino a circa 70 miliardi di punti di griglia ($2^{18}$ punti per dimensione), superando significativamente i limiti di memoria dell'implementazione di riferimento basata su cuFFT (limitata a $2^{12}$ punti).
  • Le versioni CPU e TPU hanno raggiunto rispettivamente $2^{14}$ e $2^{10}$ punti per dimensione, limitate dalla capacità di memoria.
  • Sebbene i tempi di esecuzione assoluti del metodo SFFT non fossero ancora completamente ottimizzati rispetto alle librerie cuFFT altamente ottimizzate, il comportamento di scalabilità ha indicato che l'approccio SFFT avrebbe eventualmente superato i metodi basati su FFT all'aumentare delle dimensioni del problema, in particolare per geometrie con strutture separabili dove i ranghi TT rimangono moderati.
• Accuratezza: Il metodo ha mantenuto un errore relativo inferiore al 5% per le proprietà dei materiali efficaci, controllato dal parametro di cutoff della compressione.

Significatività e Rivendicazioni
Il documento sostiene di aver gettato le basi per l'omogeneizzazione ad alte prestazioni e su larga scala basata su tensori su moderni acceleratori. Dimostra che le tecniche Tensor Train possono superare sia i colli di bottiglia di memoria che quelli computazionali nell'omogeneizzazione di dataset su scala industriale, consentendo l'omogeneizzazione di dataset massivi precedentemente impraticabili sugli acceleratori convenzionali.

Gli autori sottolineano che questo lavoro non modifica l'algoritmo SFFT fondamentale, ma si concentra sulla sua implementazione efficiente e accelerazione. Posizionano il metodo come uno strumento complementare per la modellazione multiscala guidata dai dati, capace di generare soluzioni di riferimento accurate per l'addestramento di operatori neurali. Lo studio conclude che, sebbene l'approccio sia attualmente limitato a geometrie a basso rango (ad esempio, microstrutture pixelizzate da compositi stratificati o materiali a reticolo), esso rappresenta una via percorribile verso solver quantum-inspired scalabili e basati sulla fisica per la modellazione di materiali multiscala. Gli autori rimangono modesti riguardo l'immediata applicabilità industriale per microstrutture arbitrarie, notando che sono necessari lavori futuri per affrontare la stabilità numerica sulle TPU ed estendere questi metodi a reti tensoriali di ordine superiore.