Google e Meta lançam o TorchTPU para executar o PyTorch nativamente em TPUs
O backend de código aberto, iniciado no Google em meados de 2025 com a equipe do PyTorch da Meta, já é usado em produção por alguns dos maiores clientes do Google Cloud, informou a equipe.
Ilustração de IAGoogle e Meta apresentaram o TorchTPU, um backend de código aberto que permite executar o PyTorch nativamente nos chips TPU do Google. A equipe do projeto informou que ele começou no Google em meados de 2025, em colaboração com a equipe do PyTorch na Meta, e que está em desenvolvimento ativo.
Segundo a equipe, o TorchTPU já é usado em produção, em cargas de trabalho de inferência e treinamento, por alguns dos maiores clientes do Google Cloud. A equipe observou que as TPUs estiveram historicamente vinculadas aos ecossistemas TensorFlow e JAX.
A equipe descreveu o TorchTPU como uma ruptura com tentativas anteriores. O Google lançou em 2020 o PyTorch/XLA, uma primeira versão compilada do PyTorch para TPU baseada em Lazy Tensors. Segundo a equipe, essa integração exigia mudanças intrusivas no código-fonte, sofria com severa latência de compilação, não tinha suporte ao modo de execução eager do PyTorch e ficava atrás em recursos mais recentes de computação distribuída. O TorchTPU, por sua vez, trata a execução eager como prioridade máxima e não tem dependências de JAX em seu runtime.
O objetivo declarado é que migrar uma carga de trabalho para uma TPU exija pouco mais do que trocar, no código, o dispositivo de destino de GPU para TPU. A equipe alertou que essa mudança, por si só, não entregaria o desempenho máximo da TPU, porque os formatos, o sharding e as operações de um modelo podem estar ajustados para GPUs, mas afirmou que o modelo deve rodar sem ajustes.
Até agora, segundo a equipe, o projeto:
- adicionou suporte aos modos eager e compilado, incluindo um novo modo eager chamado DeferAndFuse;
- adicionou suporte a cerca de 450 operações ATen;
- habilitou kernels personalizados escritos com JAX/Pallas e a DSL Helion;
- validou a execução sem ajustes nos 5,000 modelos mais baixados do Hugging Face;
- validou a integração com vLLM, SGLang, TorchTitan, Ray e TorchRec.
A equipe também listou limitações. Um dispositivo TPU só pode ser acessado por um processo de espaço de usuário, de modo que a colocalização de aprendizado por reforço não pode ser suportada de imediato, e a equipe disse que está investigando formas de contornar isso. Também está adicionando suporte a dinamismo limitado, porque o compilador e o runtime subjacentes presumem que os formatos dos tensores são conhecidos em tempo de compilação.
Quanto à precisão, o TorchTPU adota por padrão um modo médio que multiplica em bfloat16 e acumula em FP32. Os modos mais altos recuperam mais precisão FP32 a um custo de computação de cerca de 3 ou 6 vezes, e a equipe recomenda o padrão para a maioria das cargas de trabalho de deep learning.