Google abre o código do TorchTPU, backend nativo do PyTorch para TPUs

O que é o TorchTPU

O Google abriu o código do TorchTPU, um backend↳Back-end49 conteúdosIntegração front-end com backend: 7 decisões que evitam caos entre APIs, BFF e GraphQLDev (Back & Front) · abr 2026Como criar uma FAKE API REST para testes — JSONPlaceholderDev (Back & Front) · set 2025Construindo um aplicativo de bate-papo de IA simples com Spring AI e AngularDev (Back & Front) · jul 2025Ver tudo em Dev (Back & Front) → de PyTorch para suas TPUs (Tensor Processing Units), no repositório google-pytorch/torch_tpu no GitHub, sob licença Apache 2.0. A descrição oficial do projeto resume a proposta:

Este projeto fornece um backend PyTorch para as Tensor Processing Units (TPUs) do Google, permitindo que usuários rodem modelos e operações PyTorch em dispositivos TPU.

This project provides a PyTorch backend for Google's Tensor Processing Units (TPUs), enabling users to run PyTorch models and operations on TPU devices.Descrição do projeto, repositório google-pytorch/torch_tpu

Na prática, o projeto implementa kernels ATen, a camada de operadores que o PyTorch usa para despachar operações de tensor, voltados especificamente para TPU. É a mesma camada que os backends de CPU e CUDA já usam, o que aproxima a experiência de quem roda modelos em TPU da de quem já treina em GPU no dia a dia.

O caminho que já existia: pytorch/xla

Antes do TorchTPU, a forma estabelecida de rodar PyTorch em TPU era a biblioteca pytorch/xla, que expõe um device xla e trabalha com tensores lazy: as operações só são de fato executadas quando o código chama explicitamente xm.mark_step(). Esse modelo funciona, mas exige adaptar o loop de treino e entender uma semântica de execução diferente da eager mode que o dev já conhece de CPU e GPU.

O TorchTPU ataca esse ponto de fricção pela raiz: ao implementar os kernels diretamente na camada ATen, o objetivo é que o código PyTorch comum funcione em TPU sem reescrever o treino em torno de tensores lazy. O README lista, entre as features, integração com torch.compile() para execução em graph mode, o mesmo mecanismo de compilação que já existe no PyTorch 2.x para CPU e GPU.

XLA não desapareceu, só mudou de lugar

Aqui vale uma correção importante em relação à expectativa de que TPU nativo significa adeus ao XLA: não é bem assim. O próprio README lista entre as features um "cache de compilação para execução otimizada de computações XLA", e as instruções de build mencionam que a flag --config=wheel_common "roteia o backend XLA/MLIR compartilhado" para que ele seja compilado numa única libxla_base.so.

Em resumo: o TorchTPU muda a interface que o desenvolvedor usa (kernels ATen nativos, sem tensores lazy explícitos), mas o compilador que de fato gera código para o hardware da TPU continua sendo a pilha XLA/MLIR por baixo. É uma mudança de camada de abstração, não a eliminação do XLA do stack.

Como instalar e o que já vem pronto

A versão pública mais recente é a 0.2.0, disponível no PyPI:

pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

O projeto exige Python 3.11 ou superior, seguindo a mesma política de compatibilidade do PyTorch. Entre os recursos já documentados no README:

• Suporte a treino distribuído, com paralelismo de dados (Data Parallelism) e paralelismo de tensores (Tensor Parallelism)
• Exemplos prontos para Llama, Qwen, ResNet, DLRM, minGPT e UNet
• Cache de compilação para acelerar execuções repetidas
• Utilitários de debug e benchmarking
• Guia de usuário interativo em notebooks Marimo, ferramenta similar ao Jupyter

Os notebooks rodam com marimo edit file.py, mas precisam ser executados numa TPU VM de verdade. O próprio README recomenda uma instância v6e no GCE como caminho mais simples para quem quer testar, o que significa que não dá para só clonar o repositório e rodar localmente: é preciso provisionar hardware de TPU no Google Cloud↳Google Cloud13 conteúdosPrograma da Google Cloud no Brasil projeta formar de forma gratuita 30 mil universitáriosDev (Back & Front) · mar 2026Google Cloud OnBoard capacita estudantes e desenvolvedores de TIGestão Dev & TI · mai 2019Desenvolvedores poderão participar de treinamento gratuito do Google CloudGestão Dev & TI · mai 2019Ver tudo em DevSecOps →.

Build pela própria mão exige Bazel e paciência

Quem quer compilar o TorchTPU a partir do código fonte, em vez de instalar o wheel pronto, esbarra num sistema de build via Bazel bem mais elaborado do que o pip install sugere. O comando básico é:

bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel

A flag --config=wheel_common não é opcional, o build falha sem ela, porque é ela quem ativa o --//:wheel_build=True necessário para evitar registros duplicados na hora do import torch. Quem não tem credenciais de Remote Build Execution (RBE) precisa acrescentar --config=no_rbe depois do wheel_common, nessa ordem, porque o segundo config sobrescreve as flags de cache remoto do primeiro.

O repositório tem 4.466 commits registrados no momento da publicação, um volume que sugere desenvolvimento interno no Google por um bom tempo antes da abertura do código. Ainda assim, partes da documentação estão incompletas: a seção de dependências Python do próprio README traz só um "TBD" (a definir) no lugar da lista.

O que fica em aberto

O README não traz nenhum número de desempenho, nem comparação de custo ou velocidade contra pytorch/xla, contra GPU ou entre versões de TPU. Para quem decide entre treinar um modelo grande em TPU ou em GPU no GCP, a mudança concreta que o TorchTPU entrega até aqui é de ergonomia e esforço de engenharia, não um número de FLOPs ou de dólares por hora comprovado publicamente.

Com o projeto ainda na versão 0.2.0 e documentação de dependências marcada como pendente, faz sentido tratar o TorchTPU como uma aposta em estágio inicial: vale acompanhar os releases seguintes e testar em um projeto pequeno antes de migrar um pipeline de treino inteiro para ele. O repositório está em https://github.com/google-pytorch/torch_tpu e o guia de usuário em https://torchtpu.readthedocs.io/.

Fonte: Hacker News