TorchTPU na prática: PyTorch em TPUs e como avaliar a migração

Migrar um workload PyTorch de GPU para TPU envolve mais do que substituir o dispositivo no código. A decisão precisa preservar a qualidade do modelo, justificar o esforço de adaptação e demonstrar que a plataforma consegue provisionar, observar e recuperar a execução.

Em 20 de outubro de 2026, o Google anunciou a disponibilização do TorchTPU como código aberto, desenvolvido em colaboração com a equipe PyTorch da Meta. A proposta é executar PyTorch em Tensor Processing Units, ou TPUs, com uma experiência mais próxima das APIs usuais do framework. Isso não equivale a garantir compatibilidade ou vantagem econômica para qualquer workload. (cloud.google.com)

Este guia apresenta um processo de avaliação: primeiro compatibilidade e correção; depois desempenho; por último, operação e custo total. Os exemplos são didáticos, não benchmarks executados nem uma plataforma de produção pronta.

1. O que é TorchTPU — e o que ele não substitui

TorchTPU é um backend nativo de PyTorch, distribuído separadamente, que registra o dispositivo tpu por meio do mecanismo PrivateUse1. Ele oferece execução eager e integração com torch.compile. Não é um emulador de CUDA nem um serviço que transforma automaticamente extensões CUDA em kernels TPU. (pytorch.org)

É importante separar duas migrações:

• PyTorch/CUDA → TorchTPU: revisar dispositivos, bibliotecas, operadores e kernels específicos de GPU.
• PyTorch/XLA → TorchTPU: revisar também o modelo de execução e as integrações específicas de torch_xla.

A documentação do Google posiciona TorchTPU como substituto de PyTorch/XLA para novos workloads e como destino de migração dos existentes. Entretanto, substituir torch_xla não significa eliminar XLA: o compilador continua no caminho de execução. (docs.cloud.google.com)

Arquitetura de execução

O fluxo compilado pode ser resumido assim:

Modelo PyTorch → TorchDynamo/FX → backend TorchTPU → StableHLO → compilador XLA → runtime → TPU

ATen é a camada de operadores do PyTorch; StableHLO é uma representação intermediária de operações; XLA compila essas operações para o acelerador. O backend TPU de torch.compile utiliza XLA, não TorchInductor. (torchtpu.readthedocs.io)

Na infraestrutura, outro fluxo acontece em paralelo:

Pipeline → Job Kubernetes → agendamento em nó TPU → processo Python → runtime TPU

Para o desenho proposto neste artigo, mantenha dados, checkpoints e resultados fora do ciclo de vida do Pod. Trate o cluster como capacidade de execução, não como repositório exclusivo dos artefatos.

2. O que muda no código e no desempenho

Eager não significa ausência de compilação

A execução eager padrão preserva a experiência imperativa, mas o trabalho no dispositivo continua assíncrono e pode envolver compilação de operações. torch.compile permite otimizar regiões maiores. Leituras de tensores pela CPU e decisões Python dependentes desses valores podem forçar sincronização e interromper oportunidades de otimização. (torchtpu.readthedocs.io)

Consequência prática: meça separadamente a primeira execução, o aquecimento e o regime estável. Não use o tempo de retorno de uma chamada Python como se fosse necessariamente o tempo de conclusão no acelerador.

Shapes fazem parte do contrato operacional

Dimensões variáveis podem gerar recompilações. TorchTPU documenta dinamismo limitado por limites superiores, mas isso não significa aceitar qualquer shape sem custo adicional. (torchtpu.readthedocs.io)

Para o piloto, proponho um catálogo explícito de formatos: batch, comprimento de sequência, resolução e dtype. Compare agrupamento por tamanho, padding e entradas variáveis reais. Registre também os extremos, não apenas o caso médio.

Precisão exige validação explícita

float32 no armazenamento não garante o mesmo caminho aritmético em todos os backends. A documentação de TorchTPU descreve multiplicações matriciais baseadas em BF16, com acumulação FP32 e modos de maior precisão com custo adicional. (torchtpu.readthedocs.io)

Defina uma política de precisão para a comparação. Verifique saídas, gradientes, estabilidade da loss e qualidade final. Não escolha tolerâncias apenas para fazer o teste passar, nem exija igualdade bit a bit sem uma justificativa do domínio.

API disponível não significa caminho igualmente maduro

A matriz inicial de funcionalidades diferencia recursos validados de experimentais. Por exemplo, FSDP2 aparece como experimental, com ressalvas de desempenho, enquanto SimpleFSDP aparece entre os caminhos validados. Confirme a situação da versão adotada antes de aprovar uma arquitetura distribuída. (torchtpu.readthedocs.io)

3. Um smoke test reproduzível

Comece em uma TPU compatível, com Python e dependências alinhados à release escolhida. A documentação de instalação recomenda o build CPU de PyTorch junto ao pacote TorchTPU; não é necessário instalar o build CUDA apenas para executar em TPU. (github.com)

O comando abaixo exige que a equipe informe uma versão real previamente selecionada:

python3 -m venv .venv
source .venv/bin/activate
: "${TORCH_TPU_VERSION:?Defina uma release TorchTPU validada}"
python -m pip install "torch_tpu==${TORCH_TPU_VERSION}" --extra-index-url https://download.pytorch.org/whl/cpu
python -m pip check
python -m pip freeze > requirements-resolved.txt

Para produção, transforme a resolução aprovada em um lockfile com hashes e uma imagem imutável. pip freeze registra versões, mas não substitui sozinho esse processo.

Salve o exemplo como smoke.py. Ele mede inferência sintética em um dispositivo, com dados já residentes na TPU. A sincronização usa a API pública torch.accelerator.synchronize. (github.com)

import json
import os
import time
from importlib.metadata import version

import torch

mode = os.environ.get('EXECUTION_MODE', 'eager')
if mode not in {'eager', 'compiled'}:
    raise ValueError('EXECUTION_MODE must be eager or compiled')

torch.manual_seed(42)
device = torch.device('tpu')
model = torch.nn.Sequential(
    torch.nn.Linear(128, 256),
    torch.nn.ReLU(),
    torch.nn.Linear(256, 32),
).eval().to(device=device, dtype=torch.bfloat16)
x = torch.randn(64, 128).to(device=device, dtype=torch.bfloat16)

if mode == 'compiled':
    model = torch.compile(model, backend='tpu')

warmup = 5
iterations = 50

with torch.inference_mode():
    torch.accelerator.synchronize()
    start = time.perf_counter()
    y = model(x)
    torch.accelerator.synchronize()
    first_call_seconds = time.perf_counter() - start

    for _ in range(warmup):
        y = model(x)
    torch.accelerator.synchronize()

    start = time.perf_counter()
    for _ in range(iterations):
        y = model(x)
    torch.accelerator.synchronize()
    elapsed = time.perf_counter() - start

    finite = bool(torch.isfinite(y.float()).all().cpu().item())

if not finite:
    raise RuntimeError('Non-finite output')

print(json.dumps({
    'torch': torch.__version__,
    'torch_tpu': version('torch_tpu'),
    'mode': mode,
    'first_call_seconds': first_call_seconds,
    'mean_batch_seconds': elapsed / iterations,
    'samples_per_second': iterations * x.shape[0] / elapsed,
    'finite': finite,
}))

Execute cada modo em um processo separado:

EXECUTION_MODE=eager python smoke.py
EXECUTION_MODE=compiled python smoke.py

O programa ilustra o protocolo, não prova ganho de desempenho. A primeira chamada não inclui toda a inicialização do processo; caches persistentes também podem influenciá-la. Cinco iterações de aquecimento são apenas uma escolha didática. Amplie a janela até confirmar estabilidade e registre o estado dos caches.

O teste não cobre dataloader, backward, otimizador, comunicação distribuída ou convergência. Acrescente essas etapas antes de avaliar treinamento.

4. Kubernetes: um Job completo para o piloto

No Google Kubernetes Engine, ou GKE, o agendamento TPU utiliza seletores de acelerador e topologia, além do recurso estendido google.com/tpu. A topologia precisa corresponder ao hardware e à quantidade de chips do nó. O exemplo usa uma configuração single-host v6e 1x1. (docs.cloud.google.com)

Pré-requisitos: cluster GKE Standard compatível, node pool TPU provisionado, capacidade disponível e permissão para baixar a imagem. Construa uma imagem contendo o ambiente validado e /app/smoke.py. O endereço da imagem abaixo é ilustrativo e deve ser substituído.

apiVersion: v1
kind: Namespace
metadata:
  name: ml-evaluation
---
apiVersion: v1
kind: ServiceAccount
metadata:
  name: torchtpu-runner
  namespace: ml-evaluation
---
apiVersion: batch/v1
kind: Job
metadata:
  name: torchtpu-smoke
  namespace: ml-evaluation
spec:
  backoffLimit: 0
  activeDeadlineSeconds: 1800
  ttlSecondsAfterFinished: 86400
  template:
    metadata:
      labels:
        app: torchtpu-smoke
    spec:
      serviceAccountName: torchtpu-runner
      automountServiceAccountToken: false
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
        cloud.google.com/gke-tpu-topology: 1x1
      tolerations:
        - key: google.com/tpu
          operator: Exists
          effect: NoSchedule
      containers:
        - name: smoke
          image: us-central1-docker.pkg.dev/example-project/ml/torchtpu-smoke:validated
          imagePullPolicy: IfNotPresent
          command: [python, /app/smoke.py]
          env:
            - name: EXECUTION_MODE
              value: compiled
            - name: PYTHONUNBUFFERED
              value: '1'
          resources:
            requests:
              cpu: '2'
              memory: 8Gi
              ephemeral-storage: 1Gi
              google.com/tpu: 1
            limits:
              cpu: '4'
              memory: 16Gi
              ephemeral-storage: 4Gi
              google.com/tpu: 1

Neste desenho, backoffLimit: 0 evita mascarar falhas com novas tentativas; o deadline limita o experimento; o TTL remove o Job concluído posteriormente. Exporte logs antes dessa limpeza. Jobs são apropriados para tarefas finitas e não exigem um endpoint HTTP de readiness. (kubernetes.io)

CPU e memória são valores iniciais do exemplo, não sizing de produção. O limite de memória refere-se à RAM do contêiner, não à High Bandwidth Memory, ou HBM, da TPU. A memória do acelerador exige análise própria. (torchtpu.readthedocs.io)

Não adicione privileged: true por copiar exemplos antigos: a documentação GKE associa essa exigência a versões anteriores a 1.28. Valide permissões de dispositivo e endurecimento da imagem na combinação escolhida. (docs.cloud.google.com)

kubectl apply -f torchtpu-smoke.yaml
kubectl get pods -n ml-evaluation
kubectl logs -n ml-evaluation job/torchtpu-smoke
kubectl get events -n ml-evaluation --sort-by=.metadata.creationTimestamp

Para outra configuração, crie um Job com nome diferente. Preserve os resultados de cada execução como evidência do piloto.

5. Como avaliar a migração sem produzir um benchmark enganoso

Proponho cinco gates de aprovação:

Na comparação GPU versus TPU, fixe modelo, dataset, política de precisão, critério de qualidade e requisitos de latência. Documente diferenças de batch e paralelismo necessárias para aproveitar cada plataforma.

Para treinamento, prefira tempo e custo até a qualidade-alvo, não apenas steps por segundo. Para inferência, use throughput dentro do Service Level Objective, ou SLO, incluindo percentis de latência e falhas. Esse é o contrato de decisão recomendado para o piloto, não uma promessa de desempenho do backend.

Inclua filas, preparação dos dados, compilação, checkpoint e retomada no relatório end-to-end. Registre experimentos frios e aquecidos separadamente: TorchTPU possui caches de compilação em diferentes níveis, e o cache persistente não deve ser presumido como habilitado automaticamente. (torchtpu.readthedocs.io)

Cenário corporativo ilustrativo

Problema: um pipeline noturno gera embeddings para um catálogo, mas a equipe quer avaliar capacidade alternativa às GPUs.

Solução: manter APIs e banco de dados inalterados e deslocar apenas o worker de inferência para um piloto TPU.

Implementação: congelar uma amostra representativa, comparar qualidade de recuperação, executar as duas plataformas com o mesmo contrato de saída e publicar resultados em uma área isolada. Usar uma chave de execução para impedir duplicação na promoção dos dados.

Resultado esperado: uma decisão documentada sobre qualidade, prazo e custo. Nenhuma redução percentual é pressuposta; se a evidência for inconclusiva, mantenha o caminho atual.

6. Produção: disponibilidade, escala e recuperação

Treinamento distribuído não é replicação de uma API

TorchTPU utiliza torch.distributed com o backend tpu_dist. O planejamento deve relacionar ranks, dispositivos e comunicação à topologia física; aumentar réplicas Kubernetes isoladamente não configura o treinamento distribuído. (torchtpu.readthedocs.io)

Para a arquitetura proposta, exija recuperação coordenada do grupo, checkpoints completos e teste de falha de um worker. Salve estado do modelo, otimizador, scheduler, aleatoriedade e progresso dos dados. Defina Recovery Point Objective, ou RPO, e Recovery Time Objective, ou RTO, como critérios de aceite.

Mantenha réplicas independentes de serving em domínios de falha distintos quando houver capacidade. Não confunda um control plane regional com uma slice TPU distribuída entre zonas: são questões arquiteturais diferentes. O planejamento de topologias e capacidade deve preceder o desenho de alta disponibilidade. (docs.cloud.google.com)

Serving precisa de health checks próprios

Em um serviço HTTP, configure startupProbe para a inicialização, readinessProbe para controlar recebimento de tráfego e livenessProbe para detectar falhas recuperáveis por reinício. Probes agressivas durante compilação podem causar reinicializações indevidas. (kubernetes.io)

Como política de produção, proponho readiness apenas após carregar o modelo e concluir o warmup necessário. Dimensione capacidade com margem para falhas e upgrades; teste autoscaling com fila e latência, não somente CPU. Para treinamento, não aplique um HPA como se cada processo fosse um worker independente.

Desenvolvimento versus produção

Em upgrades, altere uma dimensão por vez: imagem, PyTorch, TorchTPU, runtime ou geração TPU. Mantenha a imagem anterior e um checkpoint compatível. Para inferência, use canário e reversão de tráfego; para treinamento, valide retomada antes de interromper o caminho anterior.

7. Segurança e observabilidade

Use uma ServiceAccount dedicada e Workload Identity Federation for GKE para acesso autorizado aos serviços Google Cloud, em vez de distribuir chaves de conta de serviço. Criar a ServiceAccount Kubernetes, sozinho, não concede permissões IAM, ou Identity and Access Management. (docs.cloud.google.com)

Como baseline, aplique RBAC, ou Role-Based Access Control, mínimo; segregação de namespaces; restrição de tráfego com NetworkPolicy; TLS; proteção de Secrets; e acesso controlado aos artefatos. Não exponha profilers ou notebooks publicamente. Valide acesso ao metadata server e aos destinos necessários antes de fechar o egress. (kubernetes.io)

Para observabilidade, proponho três painéis:

• Negócio: execuções concluídas, prazo, qualidade e custo por resultado.
• Aplicação: latência, throughput, erros, fila e tempo sem progresso.
• Infraestrutura: CPU, RAM, reinicializações, memória TPU e espera por capacidade.

Use logs estruturados com identificador da execução, revisão do modelo, versões, dtype e topologia. Defina alertas para ausência de progresso, checkpoint atrasado e deterioração de qualidade, além da indisponibilidade.

Para investigar o dispositivo, a documentação oferece integração com torch.profiler e XProf. Confira os requisitos de versão e faça capturas curtas: profiling não deve permanecer habilitado indiscriminadamente em toda execução. (torchtpu.readthedocs.io)

8. Troubleshooting orientado a evidências

As hipóteses combinam o modelo de execução TorchTPU com a investigação de recursos Kubernetes; confirme a causa antes de aplicar a correção. (torchtpu.readthedocs.io)

kubectl get nodes
kubectl get pods -n ml-evaluation -o wide
kubectl describe pod <pod> -n ml-evaluation
kubectl logs <pod> -n ml-evaluation
kubectl logs <pod> -n ml-evaluation --previous
kubectl get events -A --sort-by=.metadata.creationTimestamp

--previous recupera logs de uma instância anterior do contêiner no mesmo Pod; não recupera automaticamente logs de um Pod removido. No Job apresentado, restartPolicy: Never também limita a utilidade dessa opção. (kubernetes.io)

9. Custos, alternativas e decisão final

Use como modelo de avaliação:

custo por resultado útil = custo total atribuído à execução / resultados que atendem ao contrato de qualidade

Inclua acelerador, host, armazenamento, rede, observabilidade, preparação, compilação, ociosidade e tentativas perdidas. Amortize separadamente a engenharia de migração. Não compare apenas preço por hora ou operações teóricas por segundo.

Spot pode ser apropriado para tarefas recuperáveis, mas admite preempção. Calcule a economia líquida considerando trabalho perdido e tempo de retomada; não o trate como capacidade garantida para prazos rígidos. (docs.cloud.google.com)

Minha recomendação de triagem é:

Esses critérios são uma proposta de decisão, não um ranking universal. Evite migrar antes de resolver gargalos de dados, sem equipe para operar a nova plataforma ou sem uma justificativa econômica mensurável.

Resumo prático

• Separe compatibilidade de desempenho e desempenho de economia.
• Homologue precisão, shapes e dependências antes de escalar.
• Meça execução fria, aquecida e end-to-end.
• Diferencie treinamento distribuído de réplicas independentes de serving.
• Exija restore, rollback e capacidade disponível como gates de produção.
• Não transforme um smoke test bem-sucedido em aprovação automática da migração.

Conclusão

TorchTPU deve entrar na estratégia de Platform Engineering como uma opção de execução com contrato verificável. DevOps fornece empacotamento e promoção reproduzíveis; SRE define sinais, objetivos e recuperação; Cloud Engineering valida capacidade, isolamento e custo.

A migração está justificada quando a mesma entrega de negócio é preservada com uma combinação melhor de custo, prazo ou disponibilidade — e quando a equipe consegue explicar, medir e reverter essa mudança.