Guia do TorchTPU no PyTorch: como instalar e migrar em 2026

Migrar um projeto PyTorch de cuda para tpu pode exigir poucas mudanças no código, mas fazer tudo funcionar bem depende do acesso à TPU, da escolha correta dos pacotes, da compilação e dos formatos dos tensores.

O que o TorchTPU realmente entrega

O TorchTPU é um backend nativo do PyTorch para as TPUs do Google — não uma API de modelos hospedada nem um substituto para uma TPU VM. O anúncio do PyTorch cita execução eager, torch.compile, primitivas distribuídas e suporte a cerca de 450 operações ATen. Já o repositório público lista kernels ATen, cache de compilação, exemplos distribuídos e integração com torch.compile.

A ideia da migração é manter uma experiência familiar:

# Antes
model.to("cuda")

# Primeira adaptação para TPU
model.to("tpu")

Isso serve como teste de portabilidade, não como garantia de desempenho. O anúncio do PyTorch ressalta que o pico de performance ainda depende dos formatos do modelo, do sharding e do uso de operações adequadas ao hardware.

O que preparar antes da instalação

O TorchTPU só é útil quando há uma TPU do Google disponível. O repositório recomenda executar os notebooks em uma TPU VM e aponta uma instância v6e do Google Compute Engine como o ambiente mais simples para começar. Um laptop sem TPU consegue instalar os pacotes, mas não validar a execução na TPU.

O README público informa suporte ao Python 3.11+ e identifica o pacote torch_tpu==0.2.0. Use um ambiente isolado para evitar que o experimento herde uma versão CUDA do PyTorch ou dependências sem relação com o projeto.

Instale o torch_tpu sem puxar CUDA

Crie um ambiente virtual e instale a versão CPU do PyTorch junto com o pacote fixado do TorchTPU:

python -m venv .venv
source .venv/bin/activate

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

Segundo o README, deixar de informar o índice extra pode fazer o instalador baixar uma versão CUDA do PyTorch, que não é a dependência adequada para uma TPU VM.

Para a primeira migração, prefira o wheel publicado. Compilar a partir do código-fonte exige Bazel ou Bazelisk, além das flags documentadas pelo repositório: --config=wheel_common e, quando necessário, --config=no_rbe.

Verifique o backend com um programa PyTorch mínimo

Execute este código na TPU VM:

import torch

x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y

print(z)
print(z.device)

Se o resultado for gerado com um dispositivo tpu, isso confirma que a alocação dos tensores e a multiplicação de matrizes chegaram ao backend da TPU. Se o dispositivo não estiver disponível, confira o runtime da VM e a instalação dos pacotes antes de mexer no código do modelo.

O TorchTPU oferece execução eager para depuração e torch.compile para execução compilada. O anúncio do PyTorch relata que o modo DeferAndFuse pode ser 50% ou mais rápido que o modo eager estrito nos testes do projeto; trate esse número como uma afirmação do projeto, não como um benchmark universal.

Leve uma carga de trabalho PyTorch real para a TPU

Comece pela menor mudança possível. Só depois que o fluxo em um único dispositivo estiver validado, adicione a configuração distribuída específica da TPU.

1. Troque a seleção do dispositivo por torch.device("tpu") ou use torch.accelerator.current_accelerator() quando houver suporte.
2. Mova o modelo e os tensores de entrada para o mesmo dispositivo.
3. Execute um batch curto antes de ativar a compilação ou a execução distribuída.
4. Mantenha estáveis os tamanhos dos batches e os formatos dos tensores durante as medições; mudanças de formato podem disparar novas compilações.
5. Teste o caminho representativo do modelo e o padrão distribuído antes de assumir uma migração maior.

Um trecho distribuído de CUDA costuma usar o NCCL. No exemplo publicado pelo TorchTPU, o backend do grupo de processos para TPU é usado no lugar dele:

import torch
import torch.distributed as dist

device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)

O comando exato de inicialização distribuída depende do ambiente da TPU e da carga de trabalho. Trocar a string do dispositivo não configura automaticamente sharding, distribuição de processos ou pipelines de entrada.

Espere resultados diferentes na primeira execução e nas seguintes. O cache de compilação documentado do TorchTPU pode reutilizar grafos repetidos; por isso, registre separadamente o tempo de aquecimento e o throughput em regime estável.

Onde o TorchTPU está pronto — e onde ainda exige cuidado

O Google documenta integrações do TorchTPU com vLLM, SGLang, TorchTitan e ferramentas do Hugging Face. O anúncio sobre inferência descreve caminhos nativos para TPU no vLLM e no SGLang, incluindo movimentação do KV-cache e recursos de serving. Já o anúncio do Google Cloud relata o resultado de 52% de MFU com o Qwen 3 8B no TorchTitan e afirmações sobre uso em produção.

Na prática, as limitações são menos polidas do que sugerem as manchetes:

Quando faz sentido adotar o TorchTPU

Vale testar o TorchTPU agora se você já tem acesso a uma TPU do Google, executa modelos PyTorch e quer um caminho comum para treinamento ou inferência. O cenário mais promissor no início é o de equipes dispostas a criar perfis de desempenho e ajustar os formatos dos tensores, em vez de esperar que uma migração de uma linha entregue automaticamente a mesma eficiência de uma GPU.

Adie a migração de uma carga crítica se você precisa de uma matriz de compatibilidade madura, comportamento previsível no cold start ou de algum recurso essencial que o TorchTPU ainda não exponha. Faça a mudança em etapas: um modelo, uma fatia de TPU, um conjunto fixo de avaliação e medições separadas para tempo de instalação, primeira compilação, throughput estável e correção das saídas.

Perguntas frequentes sobre o TorchTPU

O TorchTPU é a mesma coisa que o PyTorch/XLA?

Não. O TorchTPU é um backend nativo mais recente para o PyTorch, projetado em torno da execução eager e do torch.compile. O PyTorch/XLA é o caminho mais antigo de integração com TPUs. Código existente em XLA pode exigir uma avaliação de migração separada.

Posso instalar o TorchTPU em um laptop sem TPU?

Você pode instalar os pacotes, mas não consegue validar a execução na TPU sem uma infraestrutura Google TPU compatível. A configuração prática do repositório oficial pressupõe uma TPU VM.

Por que o primeiro passo com TorchTPU foi lento?

A primeira execução pode compilar um grafo e preencher os caches. Meça as iterações seguintes separadamente e mantenha os formatos estáveis enquanto investiga novas compilações.

Trocar cuda por tpu garante o mesmo desempenho?

Não. A troca do dispositivo valida a portabilidade do código; antes de comparar desempenho, analise formatos, sharding, cobertura de operadores e tempo de aquecimento.

O TorchTPU já está pronto para produção?

O Google descreve usos de inferência e treinamento em produção, mas os componentes do ecossistema público e seus recursos têm níveis diferentes de maturidade. Teste o modelo, o mecanismo de serving e o padrão distribuído exatos antes de assumir um compromisso.

Instale o wheel fixado em uma TPU VM, execute a verificação com matrizes 2×2 e compare uma carga aquecida com o caminho atual em GPU. Assim, você separa portabilidade do código, aquecimento do compilador, cobertura de operadores e economia do hardware.