TorchTPU: Google quer tirar PyTorch da Nvidia

PyTorch preso na Nvidia é o problema que ninguém queria admitir

TorchTPU é a tentativa mais direta do Google até agora para quebrar um vício antigo da nossa área. Se você constrói com PyTorch, e a estatística diz que cerca de 85% dos apps de IA usam PyTorch, você provavelmente assume que vai rodar em GPU da Nvidia. Não é preguiça, é arquitetura. O PyTorch nasceu em 2016 dentro do Facebook, hoje Meta, numa época em que GPU da Nvidia já era padrão em pesquisa. Resultado: muita função, kernel e otimização foi pensada primeiro para CUDA, depois para o resto. Quando você tenta levar esse mesmo código para outro chip, quebra. Dá erro obscuro, cai performance, a latência estoura e o time volta para a Nvidia por puro cansaço operacional.

É nesse atrito que Google e Amazon estão mirando agora. Na terça-feira, o Google liberou o TorchTPU como projeto open source. Um dia antes, a Amazon tinha feito movimento parecido com o TorchNeuron, focado no chip Trainium e na AWS. A promessa dos dois é quase igual: mudar uma configuração e portar seu app PyTorch existente sem reescrever tudo, direto para TPU no Google Cloud ou Trainium na AWS. Parece simples no anúncio, mas quem já fez migração de acelerador sabe que o diabo mora nos detalhes de compilador, memória e serving.

O fato sem hype

O Google abriu o código do TorchTPU, que é basicamente uma camada de compatibilidade para rodar PyTorch em cima das TPUs, seus chips proprietários de IA. A ideia é permitir treinamento e inferência de modelos PyTorch no Google Cloud sem que o desenvolvedor precise reescrever o modelo. Junto veio o TPU agent suite, uma ferramenta para ajudar a migrar e depurar erros, além de uma atualização no LiteRT, que roda IA em celulares, notebooks e dispositivos embarcados, com melhor suporte a apps PyTorch.

O diferencial em relação ao anúncio da Amazon é ecossistema. O Google diz que trabalhou direto com Meta, Hugging Face, equipe do vLLM e do SGLang para garantir que bibliotecas populares funcionem no TorchTPU como funcionam no PyTorch original. Isso importa muito. Não adianta o kernel rodar se o seu pipeline de serving com vLLM ou SGLang, o seu tokenizer do Hugging Face ou o seu checkpoint não carrega direito. A Amazon, por outro lado, foi além no discurso político do open source e disse que vai submeter código ao PyTorch principal para desacoplar recursos da Nvidia, tornando o framework neutro por padrão. São duas frentes do mesmo ataque.

Como o TorchTPU funciona na visão de quem opera

Pensa no TorchTPU como um backend alternativo para o PyTorch. Hoje, quando você chama um torch.cuda ou usa autocast, compilação com torch.compile e dataloaders paralelos, por baixo existe uma pilha feita para CUDA. O que Google e Amazon fizeram foi implementar essa ponte para XLA no caso da TPU e para Neuron SDK no caso da Trainium. Na teoria, você troca o device para TPU, o compilador XLA pega o grafo e otimiza para a arquitetura sistólica da TPU, que é excelente para matmuls densas de transformers.

Na prática de operador, três pontos definem se isso presta: cobertura de operadores, compilação e memória. Cobertura porque nem toda op customizada tem kernel em TPU. Se seu modelo usa flash attention custom, kernel Triton muito específico ou extensão em C++ amarrada a CUDA, vai quebrar e vai exigir fallback para CPU, o que detona a latência. Compilação porque XLA é poderoso, mas a primeira compilação é lenta e qualquer shape dinâmico recompila tudo. Quem serve LLM com batch dinâmico e prefix caching sabe que isso pode virar um pesadelo de p99 se não travar shapes e padding direito. Memória porque TPU tem HBM generoso em pods, mas a topologia de pod exige paralelismo bem feito com sharding. Não é só trocar string de configuração.

Sobre custo e latência, dá para fazer inferência plausível mesmo sem tabela oficial final. TPU v5e e v6e no Google Cloud costumam se posicionar como alternativa mais barata por token para inferência de alto throughput quando bem utilizada, enquanto H100 segue premium e com fila. Se o TorchTPU entregar paridade com vLLM e SGLang, o ganho real aparece em workloads batch, embeddings, rerank, fine tuning de LoRA e serving de modelos abertos de 7B a 70B com alto QPS. Para protótipo pequeno e iteracão rápida, a Nvidia ainda tende a vencer pela fricção zero. Ninguém quer debugar XLA num sábado à noite para economizar 15% na conta.

O que isso muda na prática para quem constrói com IA

Quem ganha primeiro é quem tem conta de GPU estourada e escala previsível. Se você roda inferência 24 por 7 no Google Cloud ou na AWS e vive brigando por cota de H100, testar TPU ou Trainium deixou de ser projeto de três meses para virar um spike de uma semana. O caminho prático agora é isolar um endpoint não crítico, como embeddings ou um modelo de classificação, portar com TorchTPU ou TorchNeuron, medir tokens por segundo, latência p50 e p99, custo por milhão de tokens e taxa de erro de compilação. Só com esses quatro números dá para decidir.

Quem perde, no curto prazo, é quem vende conveniência em cima da Nvidia. Provedores menores de GPU cloud, wrappers de inferência e consultorias que vivem de otimizar CUDA podem ver margem comprimida se Google e Amazon realmente entregarem preço menor com fricção baixa. Para times de produto, o ajuste imediato é técnico e contratual. Técnico porque vale abstrair device, atenção e kernels customizados desde já, evitar travar código em chamadas CUDA puras e manter testes de paridade numérica. Contratual porque negociação de committed use no GCP e na AWS fica mais forte quando você tem alternativa real de chip. Uma ação prática para esta semana: pegue seu serviço PyTorch mais caro, rode o profiler, liste as ops custom e teste a importação no ambiente TPU com o agent suite para ver o que quebra. Esse relatório de dez linhas vale mais que qualquer slide sobre soberania de chips.

• Teste um workload real de inferência em TPU com vLLM ou SGLang e compare custo por milhão de tokens contra sua GPU atual.
• Mapeie kernels custom e dependências CUDA para estimar esforço real de portabilidade antes de prometer economia.
• Renegocie cotas e descontos na nuvem usando a portabilidade como alavanca, mesmo que você continue na Nvidia por enquanto.

Isso escala ou só troca o cadeado?

Aqui está a tensão que importa. Isso resolve a dependência ou só move o gargalo da Nvidia para Google e Amazon? Open source ajuda, mas TPU só roda de verdade no Google Cloud e Trainium só roda de verdade na AWS. Você sai de um lock in de hardware e entra em um lock in de nuvem, com compilador próprio, observabilidade própria e fila própria. Já vi esse filme com outros aceleradores. O dia um é lindo, o dia trinta aparece com erro de XLA, falta de exemplo no Stack Overflow e engineer da nuvem pedindo para você mudar seu modelo.

Outro ponto incômodo é manutenção. Manter paridade com o PyTorch principal dá trabalho brutal. A cada release com kernel novo, otimização nova e modelo novo, alguém precisa portar. A Amazon sinalizou que quer resolver isso na raiz, desacoplando o core da Nvidia. Se isso for aceito na Linux Foundation, é estrutural. Se ficar só em fork, vira corrida de gato e rato. Minha leitura de operador: para treinamento gigante e inferência massiva, vale o teste e pode compensar muito. Para time pequeno, com dois engenheiros de ML e pressa para entregar feature, o custo cognitivo ainda pesa mais que a economia. Portabilidade fácil no README raramente é fácil no oncall.

Conclusão prática

TorchTPU e TorchNeuron não vão derrubar a Nvidia em um trimestre, mas baixam pela primeira vez de forma séria o custo de tentar sair dela. E quando tentar fica barato, negociar fica fácil. A pergunta que fica para você é simples: quanto da sua conta de IA de hoje é performance real da Nvidia e quanto é só preguiça arquitetural que um teste de uma semana em TPU já exporia?