O Google e a Meta desenvolveram em conjunto o TorchTPU, um projeto de código aberto que adiciona um backend nativo do PyTorch às TPUs do Google. A proposta é executar cargas de trabalho PyTorch nesses aceleradores com poucas alterações no código, tanto em execução eager quanto compilada, além de oferecer suporte a computação distribuída e otimizações específicas de hardware.
A iniciativa sucede uma trajetória iniciada em 2020, quando o Google passou a viabilizar cargas PyTorch em TPUs com o PyTorch/XLA, apoiado em JAX, XLA e lazy tensors. Segundo o Google, essa abordagem tinha limitações importantes: exigia mudanças intrusivas para delimitar regiões compiláveis, apresentava alta latência de compilação, não suportava o modo imperativo eager nativo do PyTorch e demorava a acompanhar novas primitivas de computação distribuída do ecossistema.
No TorchTPU, a execução eager é uma prioridade. Para migrar uma aplicação de GPUs para TPUs, o artigo indica trocar o dispositivo cuda por tpu e o backend distribuído nccl por tpu_dist. Também apresenta a API torch.accelerator.current_accelerator, que identifica o acelerador disponível, e mostra a inicialização de um grupo distribuído com dist.init_process_group() sem especificar o backend.
O projeto disponibiliza três modos eager: strict, debug e DeferAndFuse. Este último combina execução eager e compilada. O Google afirma que ele pode ser 50% ou mais rápido que o modo Strict Eager e reduzir o pico de uso de memória HBM; trata-se de uma alegação da empresa, não de um resultado independente apresentado no artigo.
Para testar a compatibilidade com modelos existentes, o Google criou um pipeline que avalia cerca de 5.000 modelos do Hugging Face, distribuídos por mais de 500 arquiteturas. Os testes abrangem modos eager e compilado em TPUs a partir da geração v5e. A empresa descreve a metodologia como uma seleção em duas etapas, com estratificação do conjunto e equilíbrio entre adoção no mundo real e diversidade arquitetural. Segundo essa descrição, o processo evita avaliações redundantes de pequenas variações e mantém a verificação de pesos usados em produção.
Os resultados de desempenho divulgados incluem o Muse Glimmer, modelo de 30 bilhões de parâmetros implementado em PyTorch padrão. Em um único chip TPUv7, a inferência atingiu 25,7 mil tokens por segundo com comprimento de sequência de 5.120 e 22,4 mil tokens por segundo com sequência de 10.240. No treinamento, usando quatro chips TPUv7, o throughput foi de 5,6 mil e 4,6 mil tokens por segundo para essas mesmas extensões, respectivamente.
O Google também realizou experimentos de treinamento com o Qwen3 MoE 235B, escalando de 1.024 a 8.192 chips TPUv7. Para inferência de cargas MoE em grande escala com vLLM, a empresa relatou desempenho por dólar competitivo, sem apresentar no texto os números que permitiriam detalhar essa comparação.
A repercussão mencionada no artigo inclui comentários no Hacker News que elogiaram o trabalho e destacaram a riqueza do ecossistema PyTorch. A publicação também observa que, além de oferecer TPUs pelo Google Cloud, o Google as vende a clientes corporativos para implantações locais e cita a Anthropic. A possibilidade de a Meta seguir esse caminho aparece como especulação de um participante, não como anúncio confirmado.
O TorchTPU ainda está em estágio inicial e segue em evolução. O roadmap inclui suporte a formas dinâmicas, melhorias no desempenho de compilação, ampliação da cobertura de operadores, aprendizado por reforço e quantização. Esses itens são trabalhos previstos, não capacidades que o artigo afirma já estarem disponíveis.
