Análise aprofundada das TPUs

• TPU é um chip personalizado desenvolvido pelo Google para treinamento e inferência de IA em larga escala, com uma filosofia de projeto diferente da dos GPUs
• Enfatiza escalabilidade e eficiência energética, projetando em conjunto hardware (ex.: configuração de sistema em chip, grande memória on-chip) e software (compilador XLA)
• A estrutura central é baseada em array sistólico, pipelining e compilação antecipada, sendo otimizada para a maioria das operações de deep learning (especialmente multiplicação de matrizes)
• Graças às tecnologias OCI e OCS, é possível ter configuração flexível de nós, processamento paralelo de alto desempenho e escolha de diferentes topologias
• O sistema TPU se expande hierarquicamente de um único chip até multipods gigantes, permitindo treinar grandes modelos de IA e usar recursos de forma eficiente

---

Visão geral e contexto

• TPU é um chip baseado em ASIC desenvolvido pelo Google para acompanhar a expansão em larga escala dos serviços de IA
• No início, houve avaliação entre adoção de GPU, FPGA e ASIC, mas em 2013, com a ampliação do uso de deep learning em busca por voz e outros serviços, aumentou a necessidade de hardware personalizado
• Hoje, desempenha papel essencial em serviços de IA importantes como Gemini e Veo, além de ser amplamente usado em modelos de recomendação (DLRM) e outros

Estrutura de um único chip TPU

Configuração básica

• No caso do TPUv4, um único chip possui 2 TensorCores (TPUs otimizadas para inferência têm 1)
• Cada TensorCore é conectado a CMEM (128MiB) e memória HBM (32GiB)

Estrutura interna do TensorCore

• Matrix Multiply Unit (MXU) : baseada em array sistólico 128x128, dedicada à multiplicação de matrizes
• Vector Unit (VPU) : executa operações gerais elemento a elemento
• Vector Memory (VMEM; 32MiB) : copia dados da HBM e funciona como memória de preparação para o trabalho
• Scalar Unit + Scalar Memory (SMEM; 10MiB) : responsável por fluxo de controle, operações escalares e gerenciamento de endereços de memória

Diferenças estruturais em relação ao GPU

• A TPU tem memória on-chip (CMEM, VMEM, SMEM) muito maior que a de um GPU
• A capacidade de HBM é maior nos GPUs, e o número de núcleos de computação também é muito mais alto nos GPUs
• Com base no TPUv5p, são 500 TFLOPs/sec por chip e até 4.45 ExaFLOPs/sec em um pod completo (8960 chips)

Filosofia de projeto da TPU

1. Array sistólico e pipelining

• Em um array sistólico, elementos de processamento (PEs) são organizados em arranjo e passam os resultados para elementos vizinhos
• Após a entrada dos dados, é possível executar processamento contínuo sem controle adicional, e leituras/escritas em memória ocorrem apenas na entrada e na saída
• É otimizado principalmente para multiplicação de matrizes e convolução
• Com processamento em pipeline, computação e movimentação de dados ocorrem ao mesmo tempo para otimizar throughput

Desvantagem do array sistólico - esparsidade (sparsity)

• Como todos os elementos de computação ficam sempre ativos, não é adequado para matrizes esparsas
• Se no futuro os modelos de DL passarem a buscar esparsidade irregular, isso pode se tornar uma limitação

2. Compilação antecipada (AoT) e minimização de cache

• Com o codesign TPU-XLA, reduz-se ao mínimo a dependência de cache, que seria necessária para acessos irregulares à memória, diminuindo o consumo de energia
• O compilador XLA analisa previamente o grafo de operações para calcular os padrões de acesso à memória, operando com foco em memória scratchpad em vez de cache
• O @jit do JAX é uma forma intermediária entre JIT e AoT: na primeira execução, gera um grafo estático e então realiza compilação AoT com XLA
• Se o formato de entrada mudar, é necessário recompilar, e há ineficiência em padding dinâmico e repetição
• Alcança alta eficiência energética, mas tem como desvantagem menor flexibilidade

Eficiência energética do TPUv4

• Chips modernos usam memória HBM3 para reduzir o consumo de energia
• Como o gasto energético de operações de memória pode ser dezenas ou centenas de vezes maior do que o da computação, minimizar acessos à memória pode elevar bastante a eficiência

Estrutura multichip das TPUs

Nível de tray (Tray/Board; 4 chips)

• 1 tray é composto por 4 chips TPU (8 TensorCores) e um host CPU
• A conexão Host↔Chip usa PCIe, enquanto Chip↔Chip usa Inter-Core Interconnect (ICI), que oferece largura de banda mais alta

Nível de rack (Rack; 4x4x4=64 chips)

• 1 rack tem 64 chips TPU, conectados por ICI e OCS (Optical Circuit Switching) em um torus 3D (4x4x4)
• O Google diferencia rack (Rack), pod (Pod) e slice (Slice)
  • Rack: unidade física de 64 chips (= cubo)
  • Pod: maior unidade que pode ser conectada por ICI e OCS (ex.: TPUv4 = 4096 chips = 64 racks)
  • Slice: unidade abstrata de configuração arbitrária entre 4 chips e um Superpod

Vantagens do OCS

1. Wraparound : transforma cada eixo em um anel (torus 1D), reduzindo o número máximo de hops entre nós
2. Slices não contíguos flexíveis: com a estrutura de comutação do OCS, nós fisicamente separados também podem compor um único slice, facilitando uso de recursos e manutenção
3. Topologia torcida: mantendo o mesmo número de chips (x, y, z fixos), é possível alterar a estrutura de conexão para otimizar velocidade em determinados padrões de operação (ex.: twisted torus)

Casos de uso de topologia

• Formato cúbico: adequado para paralelismo de dados/tensores (largura de banda máxima)
• Formato linear (tipo charuto) : vantajoso para paralelismo em pipeline
• Torus torcido: melhora desempenho quando é necessária comunicação all-to-all (ex.: tensor parallel)

Superpod (Full Pod/Superpod; TPUv4: 4096 chips/64 racks)

• Vários racks são conectados por ICI e OCS para formar um sistema de escala gigantesca
• O desempenho de largura de banda de comunicação e processamento paralelo varia conforme o tipo de topologia do slice
• Graças ao OCS, obtém-se flexibilidade com slices não contíguos, topologias torcidas e outras opções

Nível multipod (Multi-pod/Multislice; TPUv4: mais de 4096)

• Vários pods são conectados por rede de datacenter (DCN), oferecendo infraestrutura de treinamento em larga escala, embora com largura de banda inferior à do ICI
• O treinamento do PaLM utilizou 2 pods (6144 TPUv4), com gerenciamento total de recursos em 6 pods
• No desenvolvimento de modelos grandes, o compilador XLA otimiza padrões de comunicação de acordo com as configurações definidas pelos pesquisadores (como dimensões de paralelismo)
• O XLA insere operações de comunicação entre slices e pods, permitindo treinamento distribuído em larga escala com mudanças mínimas no código

Exemplos de hardware real e diagramas

• Rack TPU: um torus 3D 4x4x4 forma uma unidade, com 2 trays (8 chips) por linha
• Tray TPUv4: na prática, há 4 portas PCIe (uma para cada TPU)
• Chip TPUv4: o ASIC fica no centro, com 4 stacks de HBM ao redor (baseado em 2 TensorCores)
• Floorplan do chip TPUv4i (para inferência): 1 TensorCore, ocupando grande área de CMEM

Encerramento

• Agradecimento ao Google TPU Research Cloud(TRC) pelo apoio à pesquisa

Referências

• TPU Multi-Slice Trainng
• Xu et al., artigo sobre GSPMD
• Jouppi et al., artigo sobre TPUv4i
• How to Scale Your Model - TPUs
• Além disso, mais de 10 artigos e materiais oficiais de apresentação