Por anos, rodar PyTorch nas TPUs do Google significava portar o código para o PyTorch/XLA ou aprender JAX. O TorchTPU, um backend nativo do PyTorch, elimina esse desvio: a TPU vira só mais um dispositivo, como uma GPU. O Hugging Face diz que suas bibliotecas Transformers, Diffusers e Kernels já rodam diretamente nele, então carregar um modelo do Hub numa TPU exige uma única configuração.
O caminho antigo tinha custos reais. O PyTorch/XLA usava tensores preguiçosos, então o código nem sempre rodava quando se esperava, e pequenas mudanças podiam disparar novas compilações do grafo difíceis de depurar. O Hugging Face criou o Optimum TPU para suavizar isso; com o TorchTPU, ele não é mais necessário para Transformers nem Diffusers. O novo backend roda em modo eager por padrão, aceita torch.compile quando se quer velocidade, lida com treinamento distribuído e traz kernels escritos em Pallas, a linguagem de kernels de baixo nível do Google. Modelos grandes demais para um chip podem ser divididos entre vários ao carregar. Num exemplo, um kernel de mistura de especialistas da biblioteca Tokamax deixou a decodificação 21,5 vezes mais rápida no Qwen3-30B-A3B e 12,8 vezes mais rápida no Gemma 4 26B-A4B com tamanho de lote 64.
O que mais importa é o custo. As TPUs são chips do próprio Google, feitos para a conta de matrizes em que as redes neurais passam o tempo, e estão disponíveis no Google Cloud de um único chip a pods com milhares. Equipes que ficaram nas GPUs da Nvidia em parte por causa do PyTorch agora podem orçar capacidade de TPU como alternativa real sem reescrever o código. Isso chega junto com um backend nativo parecido para o AWS Trainium, o TorchNeuron, e com um acordo que leva os modelos da OpenAI ao Google Cloud; juntos, apontam para um mercado menos centrado na Nvidia para ferramentas que analisamos, como PyTorch, Transformers e Diffusers.
Ainda não está pronto. O Hugging Face diz que 96,5% dos testes do Transformers que rodam em TPU passam, que já relatou 34 problemas ao projeto TorchTPU e que os testes noturnos ainda estão sendo montados. Os ganhos de velocidade citados são de kernels específicos com lote grande, não comparações completas com configurações CUDA maduras. A pergunta em aberto para quem orça uma migração é como desempenho e preço das TPUs se comparam aos das GPUs nas próprias cargas de trabalho, e isso vai precisar ser medido, não presumido.