Google TPU

Google Cloud TPUs menjalankan PyTorch secara native melalui paket torch_tpu, yang mendaftarkan perangkat tpu. torch_tpu mengompilasi operasi yang dijalankan untuk setiap bentuk input yang dilihat, sehingga menjaga bentuk input tetap statis menghindari kompilasi ulang.

Instalasi

Install torch_tpu dengan mengikuti instruksi di repo torch_tpu. torch_tpu dimuat secara otomatis saat Anda mengimpor torch. Pastikan PyTorch melihat chip TPU.

import torch

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

Ia mencetak True dan jumlah chip TPU, misalnya True 4.

Jalankan inferensi

device_map="auto" tidak mendeteksi TPU. Penempatan otomatis bergantung pada Accelerate untuk mendaftar perangkat yang tersedia, dan Accelerate tidak mendaftar perangkat TPU, sehingga model dimuat secara diam-diam di CPU. Sampaikan device_map="tpu" secara eksplisit ke from_pretrained(), atau device="tpu" ke Pipeline.

Lewatkan input generasi di CPU alih-alih memanggil .to(model.device). generate() hanya memindahkan tensor yang dikonsumsi oleh forward ke TPU. Tensor yang bertambah setiap langkah (urutan output, attention_mask) tetap di CPU, sehingga graf yang dikompilasi tetap stabil. generate mengembalikan output di perangkat yang sama dengan input. Lihat Keep generation tensors on CPU untuk detail lebih lanjut. Model menggunakan SDPA secara default di 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

Dengan cache statis, generate() mengompilasi forward pass dengan torch.compile dan backend tpu secara otomatis. Berikan disable_compile=True ke generate untuk memilih keluar.

Paralelisme tensor

torch_tpu menjalankan satu proses per chip TPU, dan setiap proses melihat chipnya sebagai tpu:0. Buat DistributedConfig dengan ukuran parallel tensor lalu teruskan ke from_pretrained(). Transformers memecah model saat dimuat.

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 memerlukan beberapa variabel lingkungan yang diatur sebelum proses dimulai, yang tidak dilakukan oleh torchrun secara mandiri. Panggil torch.tpu.distributed.set_tpu_launch_env() dan mulai torchrun dari Python.

# launch.py
import sys

import torch
from torch.distributed.run import main

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

Lewatkan argumen torchrun yang biasa ke launcher.

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

Beberapa model perlu diadaptasi sebelum dapat dipecah (sharded) di antara chip TPU.

Model yang jumlah key-value head-nya tidak habis dibagi tp_size tidak dapat melakukan sharding pada lapisan attention-nya, dan akan gagal di forward pass dengan error bentuk. Sebagai gantinya, replikasi proyeksi attention di setiap chip dengan cara menimpa aturan mereka di tp_plan. Aturan yang tidak Anda timpa akan tetap menggunakan rencana default model. Misalnya, GLM-4-9B-0414 memiliki 2 key-value head, yang tidak dapat dipecah di 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",
        },
    ),
)

Model dengan bobot terikat, seperti Qwen3-0.6B, gagal dimuat dengan tp_size lebih besar dari 1.

Gunakan kernel Hub

Perpustakaan kernels mendukung TPU sejak v0.17. Kernel TPU membungkus kode JAX dan Pallas, jadi instalasikan jax bersama kernels.

pip install "transformers[kernels]" jax

Lihat Memuat kernel untuk cara memuat kernel dari Hub.

Keterbatasan

device_map="auto" memuat model di CPU alih-alih TPU. Sampaikan device_map="tpu" secara eksplisit.
Bentuk dinamis memicu kompilasi ulang. Proses generasi melambat ketika bentuk input berubah dari satu langkah ke langkah berikutnya. Gunakan cache statis untuk generasi yang lebih cepat.
Setiap proses hanya dapat mengakses satu chip TPU. Gunakan satu proses per chip untuk menjalankan pada beberapa chip.
Integrasi Transformers untuk pelatihan dan kuantisasi di TPU masih dalam proses, sehingga beberapa fitur mungkin belum berfungsi. Silakan buat masalah jika Anda menemui kendala.

Update on GitHub