Google TPU

Google Cloud TPUs, tpu cihazını kaydeden torch_tpu paketi aracılığıyla PyTorch'u yerel olarak çalıştırır. torch_tpu, çalıştırdığı işlemleri gördüğü her giriş şekli için derler; bu nedenle giriş şekillerini statik tutmak yeniden derlemeyi önler.

Kurulum

torch_tpu deposundaki talimatları izleyerek torch_tpu paketini yükleyin. torch içe aktırıldığında torch_tpu otomatik olarak yüklenir. PyTorch'ün TPU çiplerini gördüğünü kontrol edin.

import torch

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

True ve TPU çiplerinin sayısını, örneğin True 4 olarak yazdırır.

Çıkarım çalıştır

device_map="auto" TPU’ları algılamaz. Otomatik yerleştirme, kullanılabilir cihazları listelemek için Accelerate’e dayanır ve Accelerate TPU cihazlarını listelemediğinden model sessizce CPU’ya yüklenir. Bunu önlemek için from_pretrained() çağrısına açıkça device_map="tpu" veya Pipeline çağrısına device="tpu" parametresini iletin.

Girişleri .to(model.device) çağrısı yapmak yerine CPU üzerinde iletin. generate() yalnızca forward tarafından tüketilen tensörleri TPU'ya taşır. Her adımda büyüyen tensörler (çıktı dizisi, attention_mask) CPU'da kalır, böylece derlenmiş grafik kararlı kalır. generate, çıktıyı girdilerle aynı cihazda döndürür. Daha fazla ayrıntı için Üretim tensörlerini CPU'da tutma bölümüne bakın. Modeller TPU'da varsayılan olarak SDPA kullanır.

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

Statik bir önbellekle, generate() ileri geçişi torch.compile ile ve tpu arka ucunu otomatik olarak derler. Katılmamak için disable_compile=True değerini generate fonksiyonuna iletin.

Tensor paralelliği

torch_tpu, her TPU çipi için bir süreç çalıştırır ve her süreç kendi çipini tpu:0 olarak görür. Tensor paralelliği boyutunu içeren bir DistributedConfig oluşturun ve bunu from_pretrained() fonksiyonuna iletin. Transformers, model yükleme sırasında modelin parçalarını oluşturur.

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, süreçler başlamadan önce bazı ortam değişkenlerinin ayarlanmasını gerektirir; torchrun bunu kendi başına yapmaz. torch.tpu.distributed.set_tpu_launch_env() çağrısını yapın ve torchrun’u Python’dan başlatın.

# launch.py
import sys

import torch
from torch.distributed.run import main

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

Başlatıcıya standart torchrun argümanlarını iletin.

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

Bazı modeller, TPU çipleri arasında parçalara ayrılabilmelerinden önce uyarlanmalıdır.

Key-value baş sayısı tp_size ile tam bölünemeyen modeller, dikkat katmanlarını bölüştüremez ve ileri geçişte şekil hatasıyla başarısız olur. Bunun yerine, tp_plan içindeki kurallarını geçersiz kılarak dikkat projeksiyonlarını her çipe kopyalayın. Geçersiz kılmadığınız kurallar, modelin varsayılan planını korur. Örneğin, GLM-4-9B-0414 2 key-value başına sahiptir ve bu başlar 4 çip arasında bölünemez.

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

Qwen3-0.6B gibi ağırlıkları bağlı modeller, tp_size değeri 1'den büyük olduğunda yüklenemez.

Hub çekirdeklerini kullanın

kernels kitaplığı v0.17 sürümünden itibaren TPU'yu destekler. TPU çekirdekleri JAX ve Pallas kodunu sardığı için jax paketini kernels ile birlikte yükleyin.

pip install "transformers[kernels]" jax

Hub'dan bir çekirdek yüklemek için Çekirdekleri yükleme bölümüne bakın.

Kısıtlamalar

• device_map="auto" modeli TPU yerine CPU üzerinde yükler. device_map="tpu" parametresini açıkça iletin.
• Dinamik şekiller yeniden derlemeyi tetikler. Girdi şekilleri adım adım değiştiğinde üretim yavaşlar. Daha hızlı üretim için statik bir önbellek kullanın.
• Her süreç yalnızca bir TPU çipine erişebilir. Birden fazla çipte çalıştırmak için çip başına bir süreç kullanın.
• TPU üzerinde eğitim ve kuantizasyon için Transformers entegrasyonları hâlâ devam etmektedir, bu nedenle bazı özellikler henüz çalışmayabilir. Bir sorunla karşılaşırsanız lütfen bir sorun açın.