# Google TPU

Os [Google Cloud TPUs](https://cloud.google.com/tpu) executam PyTorch nativamente por meio do pacote [torch_tpu](https://github.com/google-pytorch/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](https://github.com/google-pytorch/torch_tpu). O `torch_tpu` é carregado automaticamente quando você importa o `torch`. Verifique se o PyTorch detecta os chips TPU.

```py
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](https://huggingface.co/docs/accelerate/index) 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()](/docs/transformers/main/pt/main_classes/model#transformers.PreTrainedModel.from_pretrained), ou `device="tpu"` para [Pipeline](/docs/transformers/main/pt/main_classes/pipelines#transformers.Pipeline).

Passe os inputs de geração na CPU em vez de chamar `.to(model.device)`. [generate()](/docs/transformers/main/pt/main_classes/text_generation#transformers.GenerationMixin.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](./kv_cache#keep-generation-tensors-on-cpu) para mais detalhes. Os modelos usam SDPA por padrão no TPU.

```py
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))
```

```txt
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()](/docs/transformers/main/pt/main_classes/text_generation#transformers.GenerationMixin.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](/docs/transformers/main/pt/expert_parallelism#transformers.DistributedConfig) com o tamanho de paralelismo de tensor e passe-o para [from_pretrained()](/docs/transformers/main/pt/main_classes/model#transformers.PreTrainedModel.from_pretrained). O Transformers faz o particionamento do modelo no momento do carregamento.

```py
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.

```py
# 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.

```bash
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](https://huggingface.co/zai-org/GLM-4-9B-0414) tem 2 cabeças de chave-valor, que não podem ser divididas entre 4 chips.

```py
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](https://github.com/huggingface/kernels) oferece suporte a TPU desde a v0.17. Os kernels de TPU envolvem código JAX e [Pallas](https://docs.jax.dev/en/latest/pallas/index.html), portanto, instale `jax` junto com `kernels`.

```bash
pip install "transformers[kernels]" jax
```

Consulte [Carregamento de kernels](./kernel_doc/loading_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](https://github.com/huggingface/transformers/issues) se encontrar um problema.