O Google revelou a oitava geração de suas Unidades de Processamento Tensor (TPUs), composta por dois chips dedicados a tarefas de treino e inferência de IA.

Denominadas TPU 8t (para treino) e TPU 8i (para inferência), o Google afirmou que o hardware foi concebido em parceria com a Google DeepMind e possui "arquiteturas específicas" para suportar o treino de modelos, o desenvolvimento de agentes e cargas de trabalho de inferência.

"Nossas TPUs de oitava geração são o resultado de mais de uma década de desenvolvimento", afirmou Amin Vahdat, vice-presidente sênior e tecnólogo-chefe do Google para IA e infraestruturas. "A ideia central por trás do design original das TPU continua válida hoje em dia: ao personalizar e co-projetar o silício com hardware, redes e software, incluindo a arquitetura dos modelos e os requisitos das aplicações, conseguimos oferecer uma eficiência energética significativamente maior e um desempenho absoluto".

TPU de treino

Em uma publicação no blog que detalha as novas TPUs, Vahdat descreveu a TPU 8t como uma "potência de treino" que foi construída para "reduzir o ciclo de desenvolvimento de modelos de vanguarda de meses para semanas".

Um único superpod TPU 8t pode escalar até 9.600 chips, oferecendo dois petabytes de memória de alta largura de banda (HBM) e o dobro da largura de banda entre chips da geração anterior, a Ironwood. O Google afirmou que a arquitetura oferece 121 exaflops de desempenho de computação em FP4 para superar os estrangulamentos de largura de banda da memória, mantendo ao mesmo tempo a precisão para modelos de grande dimensão, com o desempenho de computação por pod a quase triplicar em comparação com a Ironwood.

O TPU 8t possui 19,2 Tbps de largura de banda bidirecional de expansão vertical e 400 Gbps de largura de banda de rede de expansão horizontal, tendo o Google introduzido uma nova arquitetura de rede para suportar o hardware. Denominada Virgo Network, a empresa afirmou que suporta um aumento de 4 vezes na largura de banda do centro de dados e foi construída com base em switches de alta complexidade que reduzem as camadas de rede.

Além disso, com o JAX e o Pathways, a Google afirmou que agora pode escalar para mais de 1 milhão de chips TPU num único cluster de treino, com a Virgo Network capaz de ligar mais de 134 000 chips TPU 8t com até 47 petabits por segundo de largura de banda bidirecional sem bloqueios numa única malha. Como resultado, essa estrutura oferece mais de 1,6 milhões de exaflops com um desempenho de escalabilidade quase linear, afirmou a empresa.

O Google também introduziu o TPUDirect RDMA e o TPU Direct Storage no TPU 8t. O TPUDirect RDMA permite transferências diretas de dados entre a memória e as placas de interface de rede (NICs), contornando a CPU do anfitrião e a DRAM para reduzir a latência. Entretanto, o TPU Direct Storage também contorna a CPU do anfitrião para permitir o acesso direto à memória entre o TPU e o armazenamento gerido de alta velocidade, "duplicando efetivamente a largura de banda para transferências massivas de dados", afirmou a empresa.

TPU de inferência

O TPU8i, em comparação, foi concebido para lidar com o "trabalho complexo, colaborativo e iterativo de muitos agentes especializados" que está a surgir com o advento da IA agentiva.

Concebida com maior largura de banda de memória para dar resposta a cargas de trabalho de inferência sensíveis à latência, a TPU 8i é escalável até 1152 chips num único pod. Proporciona 11,6 exaflops de desempenho de computação FP8, com uma capacidade total de HBM de 331,8 TB por pod e 19,2 Tbps de largura de banda bidirecional escalável por chip.

O Google afirmou que, no que diz respeito ao TPU 8i, também "redesenhou a pilha" para incluir quatro funcionalidades que eliminam o efeito de "sala de espera" – quando os pedidos dos usuários são intencionalmente colocados em fila ou atrasados para maximizar a utilização do hardware.

Essas incluem a combinação de 288 GB de HBM com 384 MB de SRAM integrada no chip para impedir que os processadores fiquem inativos; a duplicação do número de hosts de CPU físicos por servidor através da transição para as CPUs personalizadas Axion baseadas em Arm da Google; a duplicação da largura de banda de interligação para modelos Mixture of Experts; e a redução da latência integrada no chip em até 5 vezes com a introdução de um novo Collectives Acceleration Engine integrado no chip.

Consequentemente, Vahdat afirmou que estas inovações permitem ao TPU 8i oferecer um desempenho por dólar 80% superior ao do Ironwood.

Tanto o TPU 8t como o 8i funcionam no host de CPU baseado em Axion Arm da Google e suportam tecnologias de refrigeração líquida. A empresa afirmou que também otimizou a eficiência em toda a pilha para oferecer uma gestão de energia integrada capaz de ajustar o consumo de energia com base na procura em tempo real, resultando num desempenho por watt até duas vezes superior em comparação com o Ironwood.

"Ao controlar toda a pilha, desde o host Axion até ao acelerador, podemos otimizar a eficiência energética ao nível do sistema de formas que simplesmente não podem ser alcançadas quando o host e o chip são concebidos de forma independente", afirmou Vahdat.

Os chips estarão disponíveis ao público ainda esse ano e poderão ser utilizados como parte do AI Hypercomputer da Google – uma arquitetura de supercomputador baseada na nuvem lançada pela empresa em 2023 que combina hardware com desempenho otimizado, software aberto, estruturas de aprendizagem automática e modelos de consumo flexíveis.

*Texto traduzido e editado ao português por Bruno Faria