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

Cet article évalue les performances des opérations Tensor Train sur CPU, GPU et TPU en utilisant JAX pour adapter et accélérer un algorithme d'homogénéisation basé sur la SFFT d'inspiration quantique, permettant avec succès des simulations multi-échelles à haute résolution allant de 300 millions à 70 milliards de points de grille qui sont irréalisables avec les méthodes FFT traditionnelles basées sur GPU.

L'essentiel

Le gros problème : Trop de données, trop peu d'espace

Imaginez que vous essayiez de comprendre comment un matériau complexe (comme un alliage métallique de haute technologie ou un composite) se comporte sous l'effet d'une contrainte. Pour ce faire, les scientifiques utilisent un « microscope » pour observer la structure interne minuscule du matériau.

Par le passé, ces microscopes nous donnaient des images petites et gérables. Mais les nouvelles technologies nous donnent désormais des images à ultra-haute résolution contenant des dizaines de milliards de minuscules pixels (appelés voxels).

Le problème est que tenter d'exécuter les calculs sur ces images massives avec les méthodes traditionnelles revient à essayer de transporter une montagne de sable dans un sac en papier. L'ordinateur manque de mémoire (le sac se déchire) ou met tellement de temps à calculer que le résultat est inutile au moment où il arrive.

La solution : Une compression « inspirée du quantique »

Les auteurs proposent une nouvelle façon de gérer ces données en utilisant un tour mathématique appelé Trains de Tenseurs (Tensor Trains - TT).

Voyez les données du matériau comme un gigantesque Rubik's Cube en 3D composé de milliards de petits blocs.

• L'ancienne méthode (FFT) : Essayer de résoudre le problème en examinant chaque bloc individuellement. Cela nécessite un entrepôt massif pour stocker les données et un supercalculateur pour traiter les chiffres.
• La nouvelle méthode (Trains de Tenseurs) : Au lieu de stocker chaque bloc, vous réalisez que le cube possède un motif. Vous pouvez décrire l'ensemble en stockant seulement quelques « manuels d'instructions » (appelés cœurs) qui expliquent comment les blocs sont connectés. C'est comme compresser un film 4K en un petit fichier sans perdre l'image.

Cette méthode est qualifiée d'« inspirée du quantique » car elle emprunte une technique à la physique quantique (la Transformée de Fourier Quantique) pour résoudre les calculs, même si les auteurs l'exécutent sur des supercalculateurs classiques et non sur de véritables ordinateurs quantiques.

L'expérience : Qui est le coureur le plus rapide ?

Les auteurs voulaient voir si cette nouvelle méthode « compressée » pouvait fonctionner rapidement sur les puces informatiques modernes. Ils ont testé trois types de matériel différents :

1. CPU : Le cerveau standard d'un ordinateur (comme un bourreau de travail fiable et polyvalent).
2. GPU : Une puce conçue pour les graphismes et le traitement parallèle (comme une équipe de 10 000 fourmis travaillant ensemble).
3. TPU : Une puce spécialisée fabriquée par Google spécifiquement pour l'IA (comme une voiture de Formule 1 construite pour un circuit très spécifique).

Ils ont construit un nouveau moteur (en utilisant un outil logiciel appelé JAX) pour faire tourner leurs calculs « compressés » sur ces puces et ont chronométré leur vitesse.

Les résultats : Tout dépend de la course

L'article a révélé qu'il n'y a pas de vainqueur unique. Cela dépend de la taille du problème et du type de calcul effectué :

• Pour les tâches massives et parallèles (Le GPU gagne) : Lorsque les calculs impliquent de réaliser des millions de calculs simples en même temps (comme additionner d'immenses listes), le GPU a été le plus rapide. Il passe à l'échelle magnifiquement bien, gérant des ensembles de données massifs qui feraient planter les autres puces.
• Pour les tâches plus petites ou plus complexes (Le TPU gagne) : Pour certains types de calculs plus difficiles à diviser, le TPU s'est montré étonnamment efficace, battant souvent le CPU et parfois le GPU.
• Le CPU : C'était le plus lent, mais aussi le plus stable. Il ne plantait pas lorsque les données devenaient trop volumineuses, contrairement aux accélérateurs qui manquaient parfois de mémoire.

Un bug dans la matrice :
Les auteurs ont identé un problème spécifique avec le TPU. Lorsqu'il tentait d'effectuer un type de calcul complexe (appelé SVD) sur des nombres de très haute précision, le TPU se confondait et cessait de fonctionner correctement. Pour corriger cela, ils ont dû utiliser un « plan de secours » légèrement plus lent mais plus stable (la Décomposition Polaire) uniquement pour le TPU.

Le verdict final : Briser les limites

La partie la plus excitante de cet article est ce qu'ils ont accompli avec cette nouvelle configuration :

Ils ont réussi à exécuter des simulations d'homogénéisation sur des ensembles de données de 70 milliards de points de grille.

• Le bémol : Les meilleures méthodes traditionnelles (utilisant la FFT standard) sont tout simplement incapables de faire cela. Elles tombent en panne de mémoire bien avant d'atteindre cette taille.
• La percée : En utilisant la méthode compressée des Trains de Tenseurs sur ces accélérateurs, ils ont pu résoudre des problèmes qui étaient auparavant impossibles.

Résumé

Considérez cet article comme un essai routier d'un nouveau moteur économe en carburant (les Trains de Tenseurs) dans trois voitures différentes (CPU, GPU, TPU).

• Ils ont prouvé que ce nouveau moteur peut rouler beaucoup plus loin (gérer beaucoup plus de données) que les anciens moteurs.
• Ils ont découvert que le GPU est la meilleure voiture pour les longs trajets directs sur autoroute (données parallèles massives).
• Ils ont découvert que le TPU est excellent pour des circuits techniques spécifiques, malgré quelques particularités avec les calculs de haute précision.
• Surtout, ils ont montré qu'avec ce nouveau moteur, nous pouvons enfin traverser les « embouteillages » (ensembles de données massifs) qui étaient autrefois totalement bloqués.

Résumé technique

Résumé technique : Évaluation des performances des trains de tenseurs pour l'homogénéisation d'inspiration quantique sur les architectures TPU, GPU et CPU

Énoncé du problème
Les progrès récents de l'imagerie CT à haute résolution ont généré des ensembles de données microstructurales à ultra-haute résolution (atteignant des dizaines de milliards de voxels) qui mettent au défi les approches d'homogénéisation traditionnelles. Bien que les techniques d'homogénéisation basées sur la transformée de Fourier rapide (FFT) de pointe soient efficaces pour des ensembles de données modérés, leur empreinte mémoire et leur coût computationnel évoluent en $O(dN^d \log N)$, ce qui les rend inefficaces pour les problèmes à l'échelle industrielle. Bien que les accélérateurs matériels (GPU et TPU) offrent une puissance de calcul considérable, les exigences extrêmes en mémoire des données à haute résolution dépassent souvent leur capacité. Alors que les transformées de Fourier quantiques (QFT) offrent des accélérations exponentielles théoriques, elles restent impraticables en raison de l'absence de matériel quantique tolérant aux fautes. Par conséquent, il est nécessaire de développer des algorithmes classiques « d'inspiration quantique » qui exploitent les représentations de tenseurs de faible rang pour surmonter ces goulots d'étranglement de mémoire et de calcul.

Méthodologie
L'article étudie les performances de l'algorithme d'homogénéisation basé sur la transformée de Fourier superrapide (SFFT), qui utilise les formats Tensor Train (TT) et Tensor Train Operator (TTO) pour représenter des tenseurs d'ordre élevé. L'étude se déroule en deux phases :

1. Évaluation des opérations fondamentales : Les auteurs ont implémenté les opérations algébriques TT fondamentales (addition, multiplication, contraction, orthogonalisation et compression) en utilisant le framework JAX sur trois architectures matérielles : des CPU Intel Xeon Gold 6240R doubles, des GPU NVIDIA A100 et des Google TPU v4-8. Deux modes d'implémentation ont été comparés : un « format de liste » (les cœurs sont stockés sous forme de liste de tableaux) et un « format par lots » (les cœurs sont stockés au sein d'un seul tableau groupé). L'étude a utilisé la précision complex64 pour garantir l'exactitude, faisant fonctionner les TPU en dehors de leur régime habituel optimisé pour le BF16. La performance a été analysée via les temps d'exécution et des modèles Roofline afin de déterminer les régimes limités par la mémoire (memory-bound) ou par le calcul (compute-bound).
2. Application d'homogénéisation accélérée : Le flux de travail d'homogénéisation basé sur la SFFT a été adapté pour ces accélérateurs. Pour traiter le surcoût élevé de la compilation Just-In-Time (JIT) dans JAX lorsque les rangs des tenseurs changent dynamiquement, une stratégie de « grossissement » (coarse-graining) a été introduite. Celle-ci restreint les rangs des tenseurs à des multiples d'un rang de base ($r_0 = 16$) afin de minimiser les événements de recompilation. Pour les implémentations TPU, la compression standard basée sur la SVD a été remplacée par une compression basée sur la décomposition polaire afin d'assurer la stabilité numérique sous l'arithmétique complex64, là où la SVD a été observée comme ne convergeant pas lors de hautes discrétisations.

Principales contributions

• Premier benchmarking systématique des TPU : L'article fournit le premier benchmarking rigoureux des opérations TT fondamentales sur le matériel TPU, incluant une comparaison directe des performances avec les GPU et les CPU.
• Algèbre TT accélérée par le matériel : Il présente des implémentations efficaces de l'algèbre TT sur les accélérateurs modernes, évaluant la faisabilité du stockage en format de liste versus format par lots et identifiant des caractéristiques de performance spécifiques (ex. : comportement limité par la mémoire vs par le calcul) pour différentes opérations.
• Implémentation pratique de l'homogénéisation SFFT : Les auteurs ont adapté avec succès l'algorithme d'homogénéisation basé sur la SFFT pour l'exécution sur GPU et TPU, permettant la simulation de jeux de données allant de 300 millions à 70 milliards de points de grille — des tailles infaisables pour les implémentations de référence FFT standard basées sur GPU.
• Analyse de stabilité : Le travail identifie les instabilités numériques des opérations SVD sur TPU sous précision complex64 et propose la décomposition polaire comme alternative stable pour les régimes de haute discrétisation.

Résultats

• Performance des opérations :
  • Opérations parallèles : Pour les opérations hautement parallélisables (addition, multiplication, contraction TT-TTO), les GPU ont démontré une meilleure scalabilité aux niveaux de discrétisation élevés, finissant par surpasser les TPU. Les TPU ont montré un faible surcoût à de faibles discrétisations mais sont restés strictement limités par la mémoire sur toute la plage testée.
  • Opérations séquentielles : Pour les opérations séquentielles (orthogonalisation, compression), les TPU ont généralement surpassé les GPU sur l'ensemble du régime. Cependant, la compression basée sur la SVD sur TPU a échoué à converger autour de $2^7$ de discrétisation sous précision complex64, nécessitant le passage à la décomposition polaire.
  • Analyse Roofline : Les GPU étaient principalement limités par le calcul pour les opérations complexes, tandis que les TPU restaient limités par la mémoire pour les tâches parallèles mais transitaient vers un comportement limité par le calcul pour les tâches séquentielles à de plus grandes discrétisations.
• Scalabilité de l'homogénéisation :
  • Le solveur d'inspiration quantique basé sur GPU a réussi à monter en charge jusqu'à environ 70 milliards de points de grille ($2^{18}$ points par dimension), dépassant significativement les limites de mémoire de l'implémentation de référence basée sur cuFFT (limitée à $2^{12}$ points).
  • Les versions CPU et TPU ont atteint respectivement $2^{14}$ et $2^{10}$ points par dimension, limitées par la capacité de la mémoire.
  • Bien que les temps d'exécution absolus de la méthode SFFT ne soient pas encore totalement optimisés par rapport aux bibliothèques cuFFT hautement optimisées, le comportement de mise à l'échelle indique que l'approche SFFT finirait par surpasser les méthodes basées sur la FFT à mesure que la taille des problèmes augmente, particulièrement pour les géométries possédant des structures séparables où les rangs TT restent modérés.
• Précision : La méthode a maintenu une erreur relative inférieure à 5 % pour les propriétés matérielles effectives, contrôlée par le paramètre de coupure de la compression.

Signification et affirmations
L'article affirme établir une base pour l'homogénéisation à haute performance et à grande échelle sur les accélérateurs modernes. Il démontre que les techniques de Tensor Train peuvent surmonter à la fois les goulots d'étranglement de mémoire et de calcul dans les simulations à l'échelle industrielle, permettant l'homogénéisation de jeux de données massifs auparavant infaisables sur les accélérateurs conventionnels.

Les auteurs soulignent que ce travail ne modifie pas l'algorithme SFFT fondamental mais se concentre sur son implémentation efficace et son accélération. Ils positionnent la méthode comme un outil complémentaire pour la modélisation multi-échelle pilotée par les données, capable de générer des solutions de référence précises pour l'entraînement d'opérateurs neuronaux. L'étude conclut que, bien que l'approche soit actuellement limitée à des géométries de faible rang (ex. : microstructures pixélisées provenant de composites stratifiés ou de matériaux en réseau), elle représente une voie viable vers des solveurs d'inspiration quantique scalables et basés sur la physique pour la modélisation de matériaux multi-échelles. Les auteurs restent modestes quant à l'applicabilité industrielle immédiate pour des microstructures arbitraires, notant que des travaux futurs sont nécessaires pour traiter la stabilité numérique sur les TPU et pour étendre ces méthodes à des réseaux de tenseurs d'ordre supérieur.