# Google TPU

[Google Cloud TPUs](https://cloud.google.com/tpu) chạy PyTorch một cách bản địa thông qua gói [torch_tpu](https://github.com/google-pytorch/torch_tpu), trong đó đăng ký thiết bị `tpu`. `torch_tpu` biên dịch các thao tác mà nó thực hiện cho mỗi hình dạng đầu vào mà nó gặp, do đó việc giữ hình dạng đầu vào tĩnh giúp tránh việc biên dịch lại.

## Cài đặt

Cài đặt `torch_tpu` bằng cách làm theo các hướng dẫn trong [kho torch_tpu](https://github.com/google-pytorch/torch_tpu). `torch_tpu` sẽ được nạp tự động khi bạn import `torch`. Kiểm tra xem PyTorch có nhận thấy các chip TPU hay không.

```py
import torch

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

Nó in ra `True` và số lượng chip TPU, ví dụ `True 4`.

## Chạy suy luận

`device_map="auto"` không phát hiện được TPU. Việc đặt thiết bị tự động dựa vào [Accelerate](https://huggingface.co/docs/accelerate/index) để liệt kê các thiết bị khả dụng, và Accelerate không liệt kê các thiết bị TPU, nên mô hình sẽ được nạp lên CPU một cách im lặng. Hãy truyền `device_map="tpu"` cho [from_pretrained()](/docs/transformers/main/vi/main_classes/model#transformers.PreTrainedModel.from_pretrained) một cách rõ ràng, hoặc `device="tpu"` cho [Pipeline](/docs/transformers/main/vi/main_classes/pipelines#transformers.Pipeline).

Truyền các đầu vào sinh trên CPU thay vì gọi `.to(model.device)`. [generate()](/docs/transformers/main/vi/main_classes/text_generation#transformers.GenerationMixin.generate) chỉ di chuyển các tensor mà `forward` sử dụng sang TPU. Các tensor tăng kích thước ở mỗi bước (chuỗi đầu ra, `attention_mask`) vẫn ở trên CPU, nhờ đó đồ thị đã biên dịch được giữ ổn định. `generate` trả về đầu ra trên cùng thiết bị với đầu vào. Xem [Giữ các tensor sinh trên CPU](./kv_cache#keep-generation-tensors-on-cpu) để biết thêm chi tiết. Các mô hình sử dụng SDPA theo mặc định trên 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
```

Với bộ nhớ đệm tĩnh, [generate()](/docs/transformers/main/vi/main_classes/text_generation#transformers.GenerationMixin.generate) sẽ biên dịch bước truyền thẳng bằng `torch.compile` và backend `tpu` một cách tự động. Truyền `disable_compile=True` vào `generate` để từ chối tính năng này.

### Song song tensor

`torch_tpu` chạy một tiến trình cho mỗi chip TPU, và mỗi tiến trình nhìn thấy chip của nó là `tpu:0`. Tạo một [DistributedConfig](/docs/transformers/main/vi/expert_parallelism#transformers.DistributedConfig) với kích thước song song tensor và truyền nó vào [from_pretrained()](/docs/transformers/main/vi/main_classes/model#transformers.PreTrainedModel.from_pretrained). Transformers chia nhỏ mô hình khi tải.

```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` cần một số biến môi trường được thiết lập trước khi các tiến trình bắt đầu, điều mà `torchrun` không tự thực hiện. Hãy gọi `torch.tpu.distributed.set_tpu_launch_env()` và khởi động `torchrun` từ 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:])
```

Truyền các đối số `torchrun` thông thường cho trình khởi chạy.

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

Một số mô hình cần được điều chỉnh trước khi có thể được chia nhỏ (sharding) trên các chip TPU.

Các mô hình có số đầu key-value không chia hết cho `tp_size` sẽ không thể chia nhỏ các lớp attention và sẽ gặp lỗi shape trong bước forward. Thay vào đó, hãy nhân bản các phép chiếu attention trên mỗi chip bằng cách ghi đè các quy tắc của chúng trong `tp_plan`. Các quy tắc không được ghi đè sẽ giữ nguyên kế hoạch mặc định của mô hình. Ví dụ, [GLM-4-9B-0414](https://huggingface.co/zai-org/GLM-4-9B-0414) có 2 đầu key-value, không thể chia đều trên 4 chip.

```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",
        },
    ),
)
```

Các mô hình có trọng số liên kết, chẳng hạn như Qwen3-0.6B, không thể tải khi `tp_size` lớn hơn 1.

## Sử dụng kernel của Hub

Thư viện [kernels](https://github.com/huggingface/kernels) hỗ trợ TPU từ v0.17. Các kernel TPU bọc mã JAX và [Pallas](https://docs.jax.dev/en/latest/pallas/index.html), vì vậy hãy cài đặt `jax` cùng với `kernels`.

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

Xem [Tải kernel](./kernel_doc/loading_kernels) để biết cách tải một kernel từ Hub.

## Hạn chế

- `device_map="auto"` nạp mô hình lên CPU thay vì TPU. Hãy truyền `device_map="tpu"` một cách tường minh.
- Hình dạng động kích hoạt việc biên dịch lại. Quá trình sinh sẽ chậm lại khi hình dạng đầu vào thay đổi từ bước này sang bước khác. Hãy sử dụng bộ nhớ đệm tĩnh để sinh nhanh hơn.
- Mỗi tiến trình chỉ có thể truy cập một chip TPU. Hãy sử dụng một tiến trình cho mỗi chip để chạy trên nhiều chip.
- Các tích hợp của Transformers cho việc huấn luyện và lượng tử hóa trên TPU vẫn đang được phát triển, do đó một số tính năng có thể chưa hoạt động. Vui lòng [mở một issue](https://github.com/huggingface/transformers/issues) nếu bạn gặp sự cố.