Google TPU

Os Google Cloud TPUs executam PyTorch nativamente por meio do pacote torch_tpu, que registra o dispositivo tpu. O torch_tpu compila as operações que executa para cada forma de entrada que encontra, portanto, manter as formas de entrada estáticas evita recompilações.

Instalação

Instale o torch_tpu seguindo as instruções no repositório torch_tpu. O torch_tpu é carregado automaticamente quando você importa o torch. Verifique se o PyTorch detecta os chips TPU.

import torch

print(torch.accelerator.is_available(), torch.accelerator.device_count())

Ele imprime True e o número de chips TPU, por exemplo True 4.

Executar inferência

device_map="auto" não detecta TPUs. A colocação automática depende do Accelerate para listar os dispositivos disponíveis, e o Accelerate não lista dispositivos TPU, então o modelo é carregado silenciosamente na CPU. Passe device_map="tpu" explicitamente para from_pretrained(), ou device="tpu" para Pipeline.

Passe os inputs de geração na CPU em vez de chamar .to(model.device). generate() move apenas os tensors que forward consome para o TPU. Os tensors que crescem a cada passo (sequência de saída, attention_mask) permanecem na CPU, de modo que o gráfico compilado permanece estável. generate retorna a saída no mesmo dispositivo dos inputs. Consulte Manter os tensors de geração na CPU para mais detalhes. Os modelos usam SDPA por padrão no TPU.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B", dtype="bfloat16", device_map="tpu"
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=True, max_new_tokens=50, cache_implementation="static"
)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The French Bread Law states that a farmer can sell their harvest of bread at a certain price. The price is determined by the market price and the price of the bread, and the price is set by the farmers' decision-making committee. What are the implications of this law on

Com um cache estático, generate() compila o forward pass com torch.compile e o backend tpu automaticamente. Passe disable_compile=True para generate para desativar.

Paralelismo de tensor

torch_tpu executa um processo por chip TPU, e cada processo enxerga seu chip como tpu:0. Crie um DistributedConfig com o tamanho de paralelismo de tensor e passe-o para from_pretrained(). O Transformers faz o particionamento do modelo no momento do carregamento.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B")
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-8B",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=4),
)

# leave the inputs on CPU
inputs = tokenizer("The French Bread Law states", return_tensors="pt")
output = model.generate(
    **inputs, do_sample=False, max_new_tokens=50, cache_implementation="static"
)
if torch.distributed.get_rank() == 0:
    print(tokenizer.decode(output[0], skip_special_tokens=True))

torch_tpu requer que algumas variáveis de ambiente sejam definidas antes do início dos processos, algo que o torchrun não faz por conta própria. Chame torch.tpu.distributed.set_tpu_launch_env() e inicie o torchrun a partir do Python.

# launch.py
import sys

import torch
from torch.distributed.run import main

torch.tpu.distributed.set_tpu_launch_env()
main(sys.argv[1:])

Passe os argumentos habituais do torchrun para o launcher.

python launch.py --nproc_per_node=4 generate.py

Alguns modelos precisam de adaptações antes de poderem ser distribuídos em vários chips TPU.

Modelos cujo número de cabeças de chave-valor não é divisível por tp_size não conseguem dividir suas camadas de atenção e falham na passagem direta com um erro de forma. Em vez disso, replique as projeções de atenção em cada chip, sobrescrevendo suas regras em tp_plan. As regras que você não sobrescrever mantêm o plano padrão do modelo. Por exemplo, GLM-4-9B-0414 tem 2 cabeças de chave-valor, que não podem ser divididas entre 4 chips.

model = AutoModelForCausalLM.from_pretrained(
    "zai-org/GLM-4-9B-0414",
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(
        tp_size=4,
        tp_plan={
            "layers.*.self_attn.q_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.k_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.v_proj": "replicated_with_grad_allreduce",
            "layers.*.self_attn.o_proj": "replicated_with_grad_allreduce",
        },
    ),
)

Modelos com pesos vinculados, como o Qwen3-0.6B, falham ao carregar com tp_size maior que 1.

Use kernels do Hub

A biblioteca kernels oferece suporte a TPU desde a v0.17. Os kernels de TPU envolvem código JAX e Pallas, portanto, instale jax junto com kernels.

pip install "transformers[kernels]" jax

Consulte Carregamento de kernels para saber como carregar um kernel do Hub.

Limitações

device_map="auto" carrega o modelo na CPU em vez do TPU. Passe device_map="tpu" explicitamente.
Formas dinâmicas acionam a recompilação. A geração fica mais lenta quando as formas de entrada mudam de uma etapa para outra. Use um cache estático para uma geração mais rápida.
Cada processo pode acessar apenas um chip TPU. Use um processo por chip para executar em vários chips.
As integrações do Transformers para treinamento e quantização em TPU ainda estão em andamento, portanto, alguns recursos podem não funcionar ainda. Por favor, abra uma issue se encontrar um problema.

Update on GitHub