Google TPU

Google Cloud TPUs chạy PyTorch một cách bản địa thông qua gói 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. 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.

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 để 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 âm thầm lên CPU. Hãy truyền device_map="tpu" cho from_pretrained() một cách rõ ràng, hoặc device="tpu" cho Pipeline.

Truyền các đầu vào sinh trên CPU thay vì gọi .to(model.device). 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 để biết thêm chi tiết. Các mô hình sử dụng SDPA theo mặc định trên 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

Với bộ nhớ đệm tĩnh, 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 với kích thước song song tensor và truyền nó vào from_pretrained(). Transformers chia nhỏ mô hình khi tải.

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 đặt 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.

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

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 có 2 đầu key-value, không thể chia đều trên 4 chip.

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 hỗ trợ TPU từ v0.17. Các kernel TPU bọc mã JAX và Pallas, vì vậy hãy cài đặt jax cùng với kernels.

pip install "transformers[kernels]" jax

Xem Tải kernel để 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 nếu bạn gặp sự cố.

Update on GitHub