Google TPU

Google Cloud TPUs поддерживают PyTorch нативно через пакет torch_tpu, который регистрирует устройство tpu. torch_tpu компилирует операции для каждой формы входных данных, которую он встречает, поэтому поддержание статических форм входных данных позволяет избежать повторной компиляции.

Установка

Установите torch_tpu, следуя инструкциям в репозитории torch_tpu. torch_tpu загружается автоматически при импорте torch. Убедитесь, что PyTorch видит чипы TPU.

import torch

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

Выводится True и количество чипов TPU, например True 4.

Запуск вывода

device_map="auto" не обнаруживает TPU. Автоматическое размещение зависит от Accelerate, который перечисляет доступные устройства, но Accelerate не включает в список устройства TPU, поэтому модель молча загружается на CPU. Явно передайте device_map="tpu" в from_pretrained() или device="tpu" в Pipeline.

Передавайте входные данные для генерации на CPU, а не вызывайте .to(model.device). generate() перемещает на TPU только те тензоры, которые потребляет forward. Тензоры, растущие на каждом шаге (выходная последовательность, attention_mask), остаются на CPU, поэтому скомпилированный граф остаётся стабильным. generate возвращает результат на том же устройстве, что и входные данные. Подробнее см. в разделе Keep generation tensors on CPU. По умолчанию модели используют SDPA на 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

При использовании статического кэша generate() автоматически компилирует прямой проход с помощью torch.compile и бэкенда tpu. Чтобы отключить эту функцию, передайте disable_compile=True в generate.

Параллелизм тензоров

torch_tpu запускает один процесс на чип TPU, и каждый процесс видит свой чип как tpu:0. Создайте DistributedConfig с размером тензорного параллелизма и передайте его в from_pretrained(). Transformers разбивает модель на части при загрузке.

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 необходимо установить некоторые переменные окружения до запуска процессов, чего torchrun не делает самостоятельно. Вызовите torch.tpu.distributed.set_tpu_launch_env() и запустите torchrun из Python.

# launch.py
import sys

import torch
from torch.distributed.run import main

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

Передайте обычные аргументы torchrun в загрузчик.

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

Некоторые модели требуют адаптации перед их распределением по чипам TPU.

Модели, у которых количество ключевых и value-голов не делится нацело на tp_size, не могут шардировать слои внимания и завершаются ошибкой формы (shape error) на этапе forward pass. Вместо этого дублируйте проекции внимания на каждом чипе, переопределив их правила в tp_plan. Правила, которые вы не переопределяете, сохраняют план по умолчанию для модели. Например, у GLM-4-9B-0414 2 ключевые и value-головы, которые нельзя разделить между 4 чипами.

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, не загружаются при tp_size больше 1.

Использование ядер Hub

Библиотека kernels поддерживает TPU начиная с версии v0.17. Ядра TPU оборачивают код JAX и Pallas, поэтому установите jax вместе с kernels.

pip install "transformers[kernels]" jax

См. раздел Загрузка ядер, чтобы узнать, как загрузить ядро из Hub.

Ограничения

device_map="auto" загружает модель на CPU вместо TPU. Явно передайте device_map="tpu".
Динамические формы вызывают повторную компиляцию. Генерация замедляется, если формы входных данных меняются от шага к шагу. Используйте статический кэш для ускорения генерации.
Каждый процесс может обращаться только к одному чипу TPU. Для работы на нескольких чипах используйте по одному процессу на чип.
Интеграция Transformers для обучения и квантования на TPU всё ещё находится в разработке, поэтому некоторые функции могут пока не работать. Если вы столкнётесь с проблемой, пожалуйста, создайте issue.

Update on GitHub