Google TPU

Google Cloud TPUs 通过 torch_tpu 包原生运行 PyTorch，该包会注册 tpu 设备。torch_tpu 会为遇到的每种输入形状编译其运行的操作，因此保持输入形状静态可以避免重新编译。

安装

按照 torch_tpu 仓库中的说明安装 torch_tpu。导入 torch 时，torch_tpu 会自动加载。请检查 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 上。请显式地向 from_pretrained() 传入 device_map="tpu"，或向 Pipeline 传入 device="tpu"。

在 CPU 上传入生成输入，而不是调用 .to(model.device)。generate() 仅将 forward 所消费的张量移至 TPU。每个步骤中不断增长的张量（输出序列、attention_mask）保留在 CPU 上，从而保持编译后的图稳定。generate 返回的输出与输入位于同一设备上。有关更多详细信息，请参阅将生成张量保留在 CPU 上。在 TPU 上，模型默认使用 SDPA。

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 后端来编译前向传播。向 generate 传入 disable_compile=True 即可禁用此行为。

张量并行

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()，并从 Python 中启动 torchrun。

# 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 芯片进行分片之前，需要先进行适配。

如果模型的键值头数量不能被 tp_size 整除，则无法对注意力层进行分片，并会在前向传播中因形状错误而失败。此时，应通过覆盖 tp_plan 中的规则，将注意力投影复制到每个芯片上。未覆盖的规则将保留模型的默认计划。例如，GLM-4-9B-0414 拥有 2 个键值头，无法在 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",
        },
    ),
)

具有权重共享（tied weights）的模型（例如 Qwen3-0.6B）在 tp_size 大于 1 时无法加载。

使用 Hub 内核

kernels 库自 v0.17 起支持 TPU。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