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

Este artigo avalia o desempenho de operações de Tensor Train em CPUs, GPUs e TPUs usando JAX para adaptar e acelerar um algoritmo de homogeneização baseado em SFFT inspirado em computação quântica, permitindo com sucesso simulações multiescala de alta resolução variando de 300 milhões a 70 bilhões de pontos de grade que são inviáveis com métodos tradicionais de FFT baseados em GPU.

O essencial

O Grande Problema: Dados Demais, Espaço de Menos

Imagine que você está tentando entender como um material complexo (como uma liga metáica de alta tecnologia ou um composto) se comporta sob estresse. Para fazer isso, os cientistas usam um "microscópio" para observar a estrutura interna minúscula do material.

No passado, esses microscópios nos davam imagens pequenas e gerenciáveis. Mas a nova tecnologia agora nos fornece imagens de ultra-alta resolução contendo dezenas de bilhões de pequenos pixels (chamados de voxels).

O problema é que tentar executar os cálculos nessas imagens massivas usando métodos tradicionais é como tentar carregar uma montanha de areia em um saco de papel. O computador fica sem memória (o saco rasga) ou leva tanto tempo para calcular que o resultado se torna inútil quando chega.

A Solução: Compressão "Inspirada em Quantum"

Os autores propõem uma nova maneira de lidar com esses dados usando um truque matemático chamado Tensor Trains (TT).

Pense nos dados do material como um Cubo Mágico gigante e 3D feito de bilhões de pequenos blocos.

• O Jeito Antigo (FFT): Tentar resolver o problema olhando para cada bloco individualmente. Isso requer um armazém gigantesco para armazenar os dados e um supercomputador para processar os números.
• O Novo Jeito (Tensor Trains): Em vez de armazenar cada um dos blocos, você percebe que o cubo possui um padrão. Você pode descrever todo o conjunto armazenando apenas alguns "manuais de instrução" (chamados de cores) que dizem como os blocos se conectam. Isso é como comprimir um filme em 4K em um arquivo minúsculo sem perder a imagem.

Este método é chamado de "Inspirado em Quantum" porque toma emprestada uma técnica da física quântica (a Transformada de Fourier Quântica) para resolver a matemática, embora os autores a executem em supercomputadores comuns, não em computadores quânticos reais.

O Experimento: Quem é o Corredor Mais Rápido?

Os autores queriam ver se este novo método "comprimido" poderia rodar rápido em chips de computador modernos. Eles testaram três tipos diferentes de hardware:

1. CPU: O cérebro padrão de um computador (como um cavalo de carga confiável e de uso geral).
2. GPU: Um chip projetado para gráficos e processamento paralelo (como uma equipe de 10.000 formigas trabalhando juntas).
3. TPU: Um chip especializado feito pelo Google especificamente para IA (como um carro de Fórmula 1 construído para um tipo específico de pista).

Eles construíram um novo motor (usando uma ferramenta de software chamada JAX) para rodar sua matemática "comprimida" nesses chips e cronometrar a velocidade.

Os Resultados: Depende da Corrida

O artigo descobriu que não existe um único "vencedor". Depende do tamanho do problema e do tipo de matemática sendo feita:

• Para tarefas massivas e paralelas (A GPU vence): Quando a matemática envolve realizar milhões de cálculos simples ao mesmo tempo (como somar listas enormes), a GPU foi a mais rápida. Ela escala maravilhosamente bem, lidando com conjuntos de dados massivos que travariam os outros chips.
• Para tarefas menores ou mais complexas (A TPU vence): Para certos tipos de matemática que são mais difíceis de dividir, a TPU foi surpreendentemente eficiente, muitas vezes superando a CPU e às vezes a GPU.
• A CPU: Foi a mais lenta, mas foi a mais estável. Ela não travou quando os dados ficaram grandes demais, enquanto os aceleradores às vezes ficavam sem memória.

Uma Falha na Matrix:
Os autores encontraram um problema específico com a TPU. Ao tentar realizar um tipo específico de matemática complexa (chamada SVD) em números de alta precisão e muito grandes, a TPU ficava confusa e parava de funcionar corretamente. Para corrigir isso, eles tiveram que usar um "plano de reserva" um pouco mais lento, porém mais estável (Decomposição Polar), apenas para a TPU.

O Veredito Final: Quebrando os Limites

A parte mais emocionante do artigo é o que eles alcançaram com esta nova configuração:

Eles executaram com sucesso simulações de homogeneização em conjuntos de dados com 70 bilhões de pontos de grade.

• O Detalhe: Os melhores métodos tradicionais (usando FFT padrão) simplesmente não conseguem fazer isso. Eles ficam sem memória muito antes de atingir esse tamanho.
• O Avanço: Ao usar o método de Tensor Train "comprimido" nestes aceleradores, eles foram capazes de resolver problemas que antes eram impossíveis.

Resumo

Pense neste artigo como um teste de direção de um novo motor econômico (Tensor Trains) em três carros diferentes (CPU, GPU, TPU).

• Eles provaram que este novo motor pode dirigir muito mais longe (lidar com dados muito maiores) do que os motores antigos.
• Eles descobriram que a GPU é o melhor carro para viagens longas em estradas retas (dados paralelos massivos).
• Eles descobriram que a TPU é ótima para pistas técnicas específicas, embora tenha algumas peculiaridades com matemática de alta precisão.
• Mais importante ainda, eles mostraram que, com este novo motor, podemos finalmente atravessar "engarrafamentos" (conjuntos de dados massivos) que antes estavam completamente bloqueados.

Resumo técnico

Resumo Técnico: Benchmarking de Desempenho de Tensor Trains para Homogeneização de Inspiração Quântica em Arquiteturas TPU, GPU e CPU

Enunciado do Problema
Avanços recentes em imagens de CT de alta resolução geraram conjuntos de dados microestruturais de ultra-alta resolução (atingindo dezenas de bilhões de voxels) que desafiam as abordagens de homogeneização tradicionais. Embora as técnicas de homogeneização baseadas em Transformada Rápida de Fourier (FFT) de última geração sejam eficazes para conjuntos de dados moderados, sua pegada de memória e custo computacional escalam como $O(dN^d \log N)$, tornando-as ineficientes para problemas de escala industrial. Embora aceleradores de hardware (GPUs e TPUs) ofereçam poder computacional, os requisitos extremos de memória para dados de alta resolução frequentemente excedem sua capacidade. Embora as Transformadas de Fourier Quânticas (QFT) ofereçam acelerações exponenciais teóricas, elas permanecem impraticáveis devido à falta de hardware quântico tolerante a falhas. Consequentemente, há uma necessidade de algoritmos clássicos de "inspiração quântica" que utilizem representações de tensores de baixo posto para superar esses gargalos de memória e computação.

Metodologia
O artigo investiga o desempenho do algoritmo de homogeneização baseado em Superfast Fourier Transform (SFFT), que utiliza os formatos Tensor Train (TT) e Tensor Train Operator (TTO) para representar tensores de alta ordem. O estudo procede em duas fases:

1. Benchmarking de Operações Fundamentais: Os autores implementaram operações fundamentais de álgebra TT (adição, multiplicação, contração, ortogonalização e compressão) utilizando o framework JAX em três arquiteturas de hardware: Dual Intel Xeon Gold 6240R CPUs, NVIDIA A100 GPUs e Google TPU v4-8. Dois modos de implementação foram comparados: um modo de "formato de lista" (núcleos armazenados como uma lista de arrays) e um modo de "formato em lote" (núcleos armazenados dentro de um único array em lote). O estudo utilizou precisão complex64 para garantir a precisão, operando as TPUs fora de seu regime típico otimizado para BF16. O desempenho foi analisado via tempos de execução e modelos de Roofline para determinar os regimes limitados por memória (memory-bound) versus limitados por computação (compute-bound).
2. Aplicação de Homogeneização Acelerada: O fluxo de trabalho de homogeneização baseado em SFFT foi adaptado para esses aceleradores. Para lidar com o alto overhead de compilação Just-In-Time (JIT) no JAX quando os postos dos tensores mudam dinamicamente, uma estratégia de "granulação grossa" (coarse-graining) foi introduzida. Esta restringe os postos dos tensores a múltiplos de um posto base ($r_0 = 16$) para minimizar eventos de recompilação. Para implementações em TPU, a compressão padrão baseada em SVD foi substituída pela compressão baseada em decomposição Polar para garantir a estabilidade numérica sob aritmética complex64, onde observou-se que o SVD falhava em convergir em altas discretizações.

Principais Contribuições

• Primeiro Benchmarking Sistemático de TPU: O artigo fornece o primeiro benchmarking rigoroso de operações fundamentais de TT em hardware TPU, incluindo uma comparação direta de desempenho contra GPUs e CPUs.
• Álgebra TT Acelerada por Hardware: Apresenta implementações eficientes de álgebra TT em aceleradores modernos, avaliando a viabilidade do armazenamento em formato de lista versus formato em lote e identificando características específicas de desempenho (ex: comportamento limitado por memória vs. limitado por computação) para diferentes operações.
• Implementação Prática de Homogeneização SFFT: Os autores adaptaram com sucesso o algoritmo de homogeneização baseado em SFFT para execução em GPU e TPU, permitindo a simulação de conjuntos de dados variando de 300 milhões a 70 bilhões de pontos de grade — tamanhos inviáveis para implementações de referência de FFT baseadas em GPU padrão.
• Análise de Estabilidade: O trabalho identifica instabilidades numéricas em operações SVD em TPU sob precisão complex64 e propõe a decomposição Polar como uma alternativa estável para regimes de alta discretização.

Resultados

• Desempenho de Operação:
  • Operações Paralelas: Para operações altamente paralelizáveis (adição, multiplicação, contração TT-TTO), as GPUs demonstraram escalabilidade superior em níveis de discretização elevados, eventualmente superando as TPUs. As TPUs mostraram baixo overhead em discretizações menores, mas foram estritamente limitadas por memória em todo o intervalo testado.
  • Operações Seriais: Para operações seriais (ortogonalização, compressão), as TPUs geralmente superaram as GPUs em todo o regime. No entanto, a compressão baseada em SVD em TPUs falhou em convergir em discretizações em torno de $2^7$ sob precisão complex64, necessitando da mudança para a decomposição Polar.
  • Análise de Roofline: As GPUs foram predominantemente limitadas por computação para operações complexas, enquanto as TPUs permaneceram limitadas por memória para tarefas paralelas, mas transicionaram para um comportamento limitado por computação para tarefas seriais em maiores discretizações.
• Escalabilidade de Homogeneização:
  • O solver de inspiração quântica baseado em GPU escalou com sucesso até aproximadamente 70 bilhões de pontos de grade ($2^{18}$ pontos por dimensão), excedendo significativamente os limites de memória da implementação de referência baseada em cuFFT (limitada a $2^{12}$ pontos).
  • As versões de CPU e TPU atingiram $2^{14}$ e $2^{10}$ pontos por dimensão, respectivamente, limitadas pela capacidade de memória.
  • Embora os tempos de execução absolutos do método SFFT ainda não estivessem totalmente otimizados em comparação com bibliotecas cuFFT altamente ajustadas, o comportamento de escalonamento indicou que a abordagem SFFT eventualmente superaria os métodos baseados em FFT conforme o tamanho do problema aumentasse, particularmente para geometrias com estruturas separáveis onde os postos de TT permanecem moderados.
• Acurácia: O método manteve um erro relativo abaixo de 5% para propriedades de materiais efetivos, controlado pelo parâmetro de corte de compressão.

Significância e Alegações
O artigo afirma estabelecer uma base para homogeneização de alto desempenho e grande escala em aceleradores modernos. Demonstra que as técnicas de Tensor Train podem superar tanto os gargalos de memória quanto os de computação em simulações de escala industrial, permitindo a homogeneização de conjuntos de dados massivos anteriormente inviáveis em aceleradores convencionais.

Os autores enfatizam que este trabalho não modifica o algoritmo fundamental de SFFT, mas foca em sua implementação eficiente e aceleração. Eles posicionam o método como uma ferramenta complementar para modelagem multiescala orientada a dados, capaz de gerar soluções de referência precisas para treinamento de operadores neurais. O estudo conclui que, embora a abordagem seja atualmente limitada a geometrias de baixo posto (ex: microestruturas pixelizadas de compósitos em camadas ou materiais de rede/lattice), ela representa um caminho viável para solvers de inspiração quântica escaláveis e baseados em física para modelagem de materiais multiescala. Os autores mantêm a modéstia quanto à aplicabilidade industrial imediata para microestruturas arbitrárias, observando que trabalhos futuros são necessários para abordar a estabilidade numérica em TPUs e para estender esses métodos a redes de tensores de ordem superior.