Em 20 de outubro de 2026, o Google anunciou o TorchTPU, um backend de código aberto desenvolvido com a equipe do PyTorch da Meta para executar PyTorch nas TPUs do Google. Para começar com uma execução básica, a documentação indica instalar o pacote e selecionar o dispositivo TPU; trabalhos distribuídos também precisam de configuração específica para esse hardware.
A mudança de .to("cuda") para .to("tpu") pode bastar para iniciar uma carga de trabalho. Ela não garante, por si só, o melhor desempenho: formatos de dados, fragmentação do modelo e operações podem exigir ajustes para a arquitetura TPU.
Da execução básica ao trabalho distribuído
A documentação do Google indica pip install torch_tpu e a seleção de torch.device("tpu"). O exemplo do projeto PyTorch troca .to("cuda") por .to("tpu").
| Tipo de trabalho | Código ou backend | Configuração de execução |
| Básico | torch.device("tpu") ou .to("tpu") | Execução eager por padrão; torch.compile compila o modelo. |
| Distribuído | Backend tpu_dist, com um processo por dispositivo | Informar a topologia TPU e os dados do slice-builder ao iniciar o trabalho. |
A diferença prática aparece quando o modelo roda em vários dispositivos: além de escolher a TPU, o lançamento precisa considerar a topologia do conjunto. Para buscar o desempenho máximo, o projeto também aponta que formatos, fragmentação e operações ajustados para GPUs podem precisar de adaptação. O PyTorch relata suporte a cerca de 450 operadores ATen, com delegados estruturais ou compostos fora dessa contagem.
O que os testes publicados medem
O projeto PyTorch relatou três avaliações de compatibilidade com populações e critérios distintos. No teste Tier 1, 4.821 de 4.917 modelos selecionados por popularidade passaram sem falhas de execução, resultado apresentado como 98%. No Tier 2, 638 de 657 arquiteturas representativas alcançaram paridade numérica com uma linha de base em GPU, ou 97%. Em uma avaliação de dinamismo limitado, 834 de 839 modelos de LLM causal passaram nos testes de correção, resultado apresentado como 99%.
Essas métricas respondem a perguntas diferentes: passar sem falha não é o mesmo que atingir paridade numérica, e nenhuma delas é uma comparação padronizada de desempenho entre TPU e GPU para toda carga de trabalho.
Um exemplo concreto de adaptação envolve o Muse Glimmer 30B. Em uma implementação do TorchTitan originalmente voltada a GPU e CUDA, o projeto substituiu FlexAttention, sem suporte naquele teste, por F.scaled_dot_product_attention para executar o modelo na TPU.
O projeto também relatou um teste de treinamento do Qwen3 MoE 235B em uma configuração com 8.192 chips TPU v7, lote de 262.144 e comprimento de sequência de 4.096 tokens. Nesse teste, informou 2.104 tokens por segundo por TensorCore, utilização de 27,1% dos FLOPs do modelo (MFU) e 31,2 segundos por etapa. São números daquela configuração e daquele teste.
Reaproveitar pesos salvos
Pesos compatíveis de um modelo PyTorch treinado em GPU podem ser carregados para execução em TPU: o state_dict guarda os parâmetros do modelo, e torch.load permite remapear os tensores para o dispositivo de destino. O modelo precisa ter uma definição compatível com esses pesos.
Para continuar o treinamento, considere também o estado do otimizador e do scheduler. Eles são componentes distintos dos parâmetros armazenados no modelo e fazem parte do estado de treinamento que pode ser necessário restaurar.
