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

Este artículo evalúa el rendimiento de las operaciones de Tensor Train en CPUs, GPUs y TPUs utilizando JAX para adaptar y acelerar un algoritmo de homogeneización basado en SFFT de inspiración cuántica, permitiendo con éxito simulaciones multiescala de alta resolución que van desde 300 millones hasta 70 mil millones de puntos de malla que son inviables con los métodos tradicionales de FFT basados en GPU.

Lo esencial

El gran problema: Demasiados datos, muy poco espacio

Imagina que estás tratando de entender cómo se comporta un material complejo (como una aleación metálica de alta tecnología o un compuesto) bajo estrés. Para hacer esto, los científicos usan un "microscopio" para observar la estructura interna diminuta del material.

En el pasado, estos microscopios nos daban imágenes pequeñas y manejables. Pero la nueva tecnología ahora nos ofrece imágenes de ultra alta resolución que contienen decenas de miles de millones de diminutos píxeles (llamados vóxeles).

El problema es que intentar ejecutar los cálculos de estas imágenes masivas usando métodos tradicionales es como intentar cargar una montaña de arena en una bolsa de papel. La computadora se queda sin memoria (la bolsa se rompe) o tarda tanto en calcular que el resultado es inútil para cuando llega.

La solución: Compresión "inspirada en la cuántica"

Los autores proponen una nueva forma de manejar estos datos utilizando un truco matemático llamado Trenes de Tensores (Tensor Trains - TT).

Piensa en los datos del material como un gigantesco Cubo de Rubik 3D hecho de miles de millones de pequeños bloques.

• La forma antigua (FFT): Intentar resolver el problema mirando cada uno de los bloques individualmente. Esto requiere un almacén masivo para guardar los datos y una supercomputadora para procesar los números.
• La nueva forma (Trenes de Tensores): En lugar de guardar cada uno de los bloques, te das cuenta de que el cubo tiene un patrón. Puedes describir todo el conjunto guardando solo unos pocos "manuales de instrucciones" (llamados núcleos o cores) que te dicen cómo se conectan los bloques. Esto es como comprimir una película en 4K en un archivo diminuto sin perder la imagen.

Este método se llama "inspirado en la cuántica" porque toma prestada una técnica de la física cuántica (la Transformada de Fourier Cuántica) para resolver las matemáticas, a pesar de que los autores lo ejecutan en supercomputadoras normales, no en computadoras cuánticas reales.

El experimento: ¿Quién es el corredor más rápido?

Los autores querían ver si este nuevo método "comprimido" podía ejecutarse rápido en los chips modernos de las computadoras. Probaron tres tipos diferentes de hardware:

1. CPU: El cerebro estándar de una computadora (como un caballo de trabajo confiable y de uso general).
2. GPU: Un chip diseñado para gráficos y procesamiento paralelo (como un equipo de 10,000 hormigas trabajando juntas).
3. TPU: Un chip especializado fabricado por Google específicamente para la IA (como un auto de Fórmula 1 construido para un tipo de pista específico).

Construyeron un nuevo motor (usando una herramienta de software llamada JAX) para ejecutar sus matemáticas "comprimidas" en estos chips y cronometraron qué tan rápido iban.

Los resultados: Depende de la carrera

El artículo encontró que no hay un único "ganador". Todo depende del tamaño del problema y del tipo de matemática que se esté realizando:

• Para tareas masivas y paralelas (La GPU gana): Cuando las matemáticas implican realizar millones de cálculos simples a la vez (como sumar listas enormes), la GPU fue la más rápida. Escala de manera excelente, manejando conjuntos de datos masivos que colapsarían a los otros chips.
• Para tareas más pequeñas o complejas (La TPU gana): Para ciertos tipos de matemáticas que son más difíciles de dividir, la TPU fue sorprendentemente eficiente, superando a menudo a la CPU y, a veces, a la GPU.
• La CPU: Fue la más lenta, pero la más estable. No colapsó cuando los datos se volvieron demasiado grandes, mientras que los aceleradores a veces se quedaron sin memoria.

Un fallo en la Matrix:
Los autores encontraron un problema específico con la TPU. Al intentar realizar un tipo específico de matemática compleja (llamada SVD) con números de alta precisión muy grandes, la TPU se confundía y dejaba de funcionar correctamente. Para solucionar esto, tuvieron que usar un "plan de respaldo" ligeramente más lento pero más estable (Descomposición Polar) solo para la TPU.

El veredicto final: Rompiendo los límites

La parte más emocionante del artículo es lo que lograron con esta nueva configuración:

Lograron ejecutar simulaciones de homogeneización en conjuntos de datos con 70 mil millones de puntos de malla.

• El detalle: Los mejores métodos tradicionales (usando la FFT estándar) simplemente no pueden hacer esto. Se quedan sin memoria mucho antes de alcanzar ese tamaño.
• El avance: Al usar el método de Trenes de Tensores "comprimido" en estos aceleradores, pudieron resolver problemas que antes eran imposibles.

Resumen

Piensa en este artículo como una prueba de manejo para un nuevo motor de bajo consumo (Trenes de Tensores) en tres autos diferentes (CPU, GPU, TPU).

• Demostraron que este motor puede conducir mucho más lejos (manejar datos mucho más grandes) que los motores antiguos.
• Encontraron que la GPU es el mejor auto para viajes largos por autopistas rectas (datos masivos en paralelo).
• Encontraron que la TPU es excelente para pistas técnicas específicas, aunque tiene algunas peculiaridades con las matemáticas de alta precisión.
• Lo más importante es que demostraron que, con este nuevo motor, finalmente podemos atravesar "atascos de tráfico" (conjuntos de datos masivos) que antes estaban completamente bloqueados.

Resumen técnico

Resumen Técnico: Evaluación del Rendimiento de los Trenes de Tensores para la Homogeneización de Inspiración Cuántica en Arquitecturas TPU, GPU y CPU

Planteamiento del Problema
Los avances recientes en la tomografía computarizada (CT) de alta resolución han generado conjuntos de datos microestructurales de ultra alta resolución (que alcanzan decenas de miles de millones de vóxeles) que desafían los enfoques de homogeneización tradicionales. Aunque las técnicas de homogeneización basadas en la Transformada Rápida de Fourier (FFT) de vanguardia son efectivas para conjuntos de datos moderados, su huella de memoria y costo computacional escalan como $O(dN^d \log N)$, lo que las hace ineficientes para problemas de escala industrial. Si bien los aceleradores de hardware (GPUs y TPUs) ofrecen potencia computacional, los extremos requisitos de memoria para datos de alta resolución suelen exceder su capacidad. Aunque las Transformadas de Fourier Cuánticas (QFT) ofrecen aceleraciones exponenciales teóricas, siguen siendo impracticables debido a la falta de hardware cuántico tolerante a fallos. En consecuencia, existe la necesidad de algoritmos clásicos de "inspiración cuántica" que aprovechen las representaciones de tensores de bajo rango para superar estos cuellos de botella de memoria y computación.

Metodología
El artículo investiga el rendimiento del algoritmo de homogeneización basado en la Transformada de Fourier Superrápida (SFFT), el cual utiliza los formatos de Tren de Tensores (TT) y Operador de Tren de Tensores (TTO) para representar tensores de alto orden. El estudio se desarrolla en dos fases:

1. Evaluación de Operaciones Fundamentales: Los autores implementaron operaciones algebraicas fundamentales de TT (suma, multiplicación, contracción, ortogonalización y compresión) utilizando el marco de trabajo JAX a través de tres arquitecturas de hardware: CPUs duales Intel Xeon Gold 6240R, GPUs NVIDIA A100 y Google TPU v4-8. Se compararon dos modos de implementación: un "formato de lista" (núcleos almacenados como una lista de matrices) y un "formato por lotes" (núcleos almacenados dentro de una única matriz por lotes). El estudio utilizó precisión complex64 para asegurar la exactitud, operando las TPUs fuera de su régimen típico optimizado para BF16. El rendimiento se analizó mediante tiempos de ejecución y modelos Roofline para determinar los regímenes limitados por memoria o por cómputo.
2. Aplicación de Homogeneización Acelerada: El flujo de trabajo de homogeneización basado en SFFT fue adaptado para estos aceleradores. Para abordar la alta sobrecarga de la compilación Just-In-Time (JIT) en JAX cuando los rangos de los tensores cambian dinámicamente, se introdujo una estrategia de "granularidad gruesa". Esta restringe los rangos de los tensores a múltiplos de un rango base ($r_0 = 16$) para minimizar los eventos de recompilación. Para las implementaciones en TPU, la compresión estándar basada en SVD fue reemplazada por la compresión basada en descomposición polar para asegurar la estabilidad numérica bajo aritmética complex64, donde se observó que la SVD fallaba en converger en regímenes de alta discretización.

Contribuciones Clave

• Primera Evaluación Sistemática en TPU: El artículo proporciona la primera evaluación rigurosa de las operaciones fundamentales de TT en hardware TPU, incluyendo una comparación directa de rendimiento contra GPUs y CPUs.
• Álgebra de TT Acelerada por Hardware: Presenta implementaciones eficientes de álgebra TT en aceleradores modernos, evaluando la viabilidad del almacenamiento en formato de lista frente al formato por lotes e identificando características de rendimiento específicas (por ejemplo, comportamiento limitado por memoria frente a limitado por cómputo) para diferentes operaciones.
• Implementación Práctica de la Homogeneización SFFT: Los autores adaptaron con éxito el algoritmo de homogeneización basado en SFFT para la ejecución en GPU y TPU, permitiendo la simulación de conjuntos de datos que van desde 300 millones hasta 70 mil millones de puntos de rejilla—tamaños que son inviables para las implementaciones de referencia de FFT basadas en GPU estándar.
• Análisis de Estabilidad: El trabajo identifica inestabilidades numéricas en las operaciones SVD de TPU bajo precisión complex64 y propone la descomposición polar como una alternativa estable para regímenes de alta discretización.

Resultados

• Rendimiento de las Operaciones:
  • Operaciones Paralelas: Para operaciones altamente paralelizables (suma, multiplicación, contracción TT-TTO), las GPUs demostraron una escalabilidad superior en niveles de discretización altos, superando eventualmente a las TPUs. Las TPUs mostraron una baja sobrecarga en discretizaciones bajas, pero fueron estrictamente limitadas por memoria en todo el rango probado.
  • Operaciones Seriales: Para operaciones seriales (ortogonalización, compresión), las TPUs generalmente superaron a las GPUs en todo el régimen. Sin embargo, la compresión basada en SVD en TPUs falló en converger en discretizaciones alrededor de $2^7$ bajo precisión complex64, lo que obligó al cambio a la descomposición polar.
  • Análisis Roofline: Las GPUs fueron predominantemente limitadas por cómputo para operaciones complejas, mientras que las TPUs permanecieron limitadas por memoria para tareas paralelas, pero transicionaron hacia un comportamiento limitado por cómputo para tareas seriales en discretizaciones mayores.
• Escalabilidad de la Homogeneización:
  • El solver de inspiración cuántica basado en GPU escaló con éxito hasta aproximadamente 70 mil millones de puntos de rejilla ($2^{18}$ puntos por dimensión), excediendo significativamente los límites de memoria de la implementación de referencia basada en cuFFT (limitada a $2^{12}$ puntos).
  • Las versiones de CPU y TPU alcanzaron $2^{14}$ y $2^{10}$ puntos por dimensión, respectivamente, limitadas por la capacidad de memoria.
  • Aunque los tiempos de ejecución absolutos del método SFFT aún no están totalmente optimizados en comparación con las bibliotecas cuFFT altamente ajustadas, el comportamiento de escalado indicó que el enfoque SFFT eventualmente superaría a los métodos basados en FFT a medida que aumentaran los tamaños de los problemas, particularmente para geometrías con estructuras separables donde los rangos TT se mantienen moderados.
• Exactitud: El método mantuvo un error relativo inferior al 5% para las propiedades de materiales efectivas, controlado por el parámetro de corte de compresión.

Significancia y Reivindicaciones
El artículo afirma establecer una base para la homogeneización de alto rendimiento y gran escala en aceleradores modernos. Demuestra que las técnicas de Tren de Tensores pueden superar tanto los cuellos de botella de memoria como los de computación en simulaciones de escala industrial, permitiendo la homogeneización de conjuntos de datos masivos previamente inviables en aceleradores convencionales.

Los autores enfatizan que este trabajo no modifica el algoritmo fundamental de SFFT, sino que se centra en su implementación eficiente y aceleración. Posicionan el método como una herramienta complementaria para el modelado multiescala basado en datos, capaz de generar soluciones de referencia precisas para entrenar operadores neuronales. El estudio concluye que, si bien el enfoque está actualmente limitado a geometrías de bajo rango (por ejemplo, microestructuras pixeladas de compuestos laminados o materiales de red), representa una vía viable hacia solvers de inspiración cuántica escalables y basados en la física para el modelado de materiales multiescala. Los autores mantienen la modestia respecto a la aplicabilidad industrial inmediata para microestructuras arbitrarias, señalando que se requiere trabajo futuro para abordar la estabilidad numérica en TPUs y extender estos métodos a redes de tensores de orden superior.