如果你曾经尝试在 TPU 上运行 PyTorch，你很可能将代码移植到了 torch_xla，或者最终不得不学习 JAX。有了 TorchTPU 这个原生的 PyTorch 后端，TPU 现在只是另一种设备。它开启了 Hugging Face 生态系统对 TPU 的原生支持。Transformers、Diffusers 和 Kernels 可在 TPU 上原生运行。从 Hub 选择一个模型，将其发送到“tpu”，就完成了。

等等，TPU 是什么？🤔

TPU 代表张量处理单元。TPU 是 Google 自有的 AI 加速器。它们是专为一项工作而打造的芯片：神经网络大部分时间所进行的矩阵运算。Google 已经设计它们大约十年，并用它们来训练和提供自己的模型，包括 Gemini。

TPU 可在 Google Cloud 上使用，从单个芯片到数千芯片的 pod。当你需要大量计算时，无论是训练还是大规模服务，它们都表现出色。

一直以来的问题都是软件。TPU 软件栈是围绕 TensorFlow 和 JAX 发展起来的，PyTorch 用户要进入其中就得更加努力。

旧有方式

直到现在，在 TPU 上运行 PyTorch 意味着要通过 PyTorch/XLA（torch_xla）。它能工作，但有其怪癖。张量是惰性的（是张量惰性，不是你），所以你的代码并不总是在你以为的时候运行。小改动可能触发图重新编译，调试也很难理清思路。

为了让生活更轻松，我们构建了 Optimum TPU。它为 Hugging Face 模型提供了在 TPU 上的推理和训练支持。现在 TorchTPU 带来了原生支持，运行带有 Transformers 和 Diffusers 的 Hugging Face 模型已不再需要 Optimum TPU。

进入 TorchTPU 🚀

TorchTPU 为 PyTorch 添加了 TPU 作为新后端。你编写常规的 PyTorch 代码。它自带针对 TPU 调优的 Pallas 内核、分布式训练支持，以及当你想要更高速度时的 torch.compile 集成。在 TPU VM 上安装只需一条命令：

Shell
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu

额外的索引会给你 torch 的 CPU 构建版本。没有它，pip 会拉取大得多的 CUDA 构建版本，而在 TPU 上你并不需要它。

你最爱的库，现在支持 TPU

Transformers

Transformers v5 围绕一个理念构建：成为模型定义的真理之源，使它们可以在任何地方运行。“任何地方”现在包括 TPU。要在 TPU 上加载模型，传入 device_map="tpu"。就是这样。

Python
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
 model_id,
 device_map="tpu", # will load model directly on the TPU
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The future of AI is", return_tensors="pt").to("tpu")
out = model.generate(**inputs, max_new_tokens=64)
print(tokenizer.decode(out[0], skip_special_tokens=True))

模型默认在 eager 模式下运行。你也可以使用 torch.compile 以获得更好性能，或在分布式设置中运行。

图像？视频？Diffusers 在 TPU 上绘制它们 🎨

TPU 现在是 Diffusers 中的一等设备。对于原型设计或调试，eager 模式只需 pipe.to("tpu")。对于部署，编译模式只需添加一次编译调用。你的流水线和代码保持不变，也没有需要维护的 TPU 分支。

Python
import torch
from diffusers import StableDiffusion3Pipeline

pipe = StableDiffusion3Pipeline.from_pretrained(
 "stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16
)
pipe.to("tpu") # eager mode: that's the whole change
# optional: compiled mode
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)

image = pipe("a corgi astronaut planting a flag on a pizza moon").images[0]

第一次调用会进行编译，之后相同图像尺寸的每次调用都会复用该图。

模型太大，一个芯片装不下？向 from_pretrained 传入 TensorParallelConfig，transformer 在加载时就会跨芯片分片，因此完整模型永远不必装进一个地方。许多流行流水线开箱即用支持它，不支持的那些也只需一个 tp_plan。Diffusers TPU 指南有完整配方。

以下是一些在 Cloud TPU v6e-4 上运行的示例：

代码
TP=4 (each model split across four chips) · bf16 · batch 1

MiniMax-H3（视频 + 音频） Qwen-Image-2.1

“图片中毛茸茸的 Hugging Face 吉祥物，一个圆圆的黄色笑脸，带着红扑扑的脸颊和两只张开的小手，站在一艘古希腊三层桨战船的船头……海浪拍打声、木材吱嘎声、绳索中的风声、远处的海鸥声，以及低沉的史诗管弦乐渐强。” · 1344×768，5.2 秒，49 步 · 编译后 188 秒，eager 262 秒 “一只戴着巫师帽的水豚，在烛光下读书，油画” · 1024×1024，40 步 · 编译后 3.5 秒，eager 8.0 秒 “一个霓虹店招写着‘QWEN IMAGE 2.1’，雨夜，湿路面上的倒影” · 2048×2048，40 步

Kernels

Kernels Hub 提供即用型 TPU 内核。用 Kernels 库加载一个，给你的模型加速。

TPU 内核用 Pallas 编写，让你对硬件有低级控制。你也可以引入 Tokamax 库中的内核。以下是将 Tokamax 的 ragged-dot 内核插入混合专家模型的示例：

Python
from kernels import get_kernel
from transformers import AutoModelForCausalLM
from transformers.integrations.moe import ALL_EXPERTS_FUNCTIONS

repo_id = "tengomucho/ragged-dot-tpu"
kernel = get_kernel(repo_id, version=1, trust_remote_code=[repo_id])
ALL_EXPERTS_FUNCTIONS.register("tokamax_ragged_dot", kernel.experts_forward)
model = AutoModelForCausalLM.from_pretrained(
 "Qwen/Qwen3-30B-A3B-Instruct-2507",
 experts_implementation="tokamax_ragged_dot",
)

批次越大，收益越大。在批次大小 64 时，Qwen3-30B-A3B 的解码速度提升 21.5 倍，Gemma 4 26B-A4B 提升 12.8 倍。

在 TPU v6e-8 上的解码吞吐量，Hub 内核与默认实现对比。

我们测试它，这样你就不必自己测了 🧪

我们通过运行库测试套件持续改进 transformers 中的 TPU 支持。我们正在努力添加 TPU 上的夜间 CI，并在（transformers dashboard）[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard] 中提供结果。

Transformers CI 仪表板，TPU 结果与其他加速器并列。进行中！

还不是全部绿灯，我们对此没问题。在 TPU 上运行的测试中，96.5% 通过。我们发现的每个失败都会变成一个小型复现脚本和 TorchTPU 仓库上的 issue。到目前为止我们已报告 34 个 issue，其中 19 个已修复。在 TPU 上运行还捕获了 Transformers 本身的一些 bug，例如 StaticCache 和 Qwen3 张量并行计划中的问题。我们已在上游修复了这些。

如果有什么在 TPU 上出问题，我们希望在你之前发现。

开始使用

对 PyTorch 用户来说，TPU 不再是支线任务。安装 TorchTPU，从 Hub 抓取一个模型，然后发送到“tpu”。

GitHub 上的 TorchTPU

Hub 上的 TPU 内核

Cloud TPU 文档

我们迫不及待想看到你构建的成果 🤗。如果遇到问题，请提交 issue 告诉我们。