如果你曾经尝试在 TPU 上运行 PyTorch，你可能把代码移植到了torch_xla，或者最终不得不去学 JAX。有了TorchTPU这个原生 PyTorch 后端，TPU 现在只是另一种设备。它为 HF Mirror 生态带来了对 TPU 的原生支持。Transformers、Diffusers 和 Kernels 都能在 TPU 上原生运行。从 Hub 里挑一个模型，把它发送到"tpu"，就完成了。
等等，TPU 是什么？🤔
TPU 代表张量处理单元（Tensor Processing Unit）。TPU 是 Google 自家的 AI 加速器。这些芯片专为一件事而生：神经网络大部分时间都在做的矩阵运算。Google 设计它们已有约十年，并用它们来训练和部署自家模型，包括 Gemini。
TPU 可在Google Cloud上获取，从单颗芯片到由数千颗芯片组成的 pod 都有。当你需要大量算力时——无论是训练还是大规模部署——它们就大放异彩。
问题一直出在软件上。TPU 技术栈是围绕 TensorFlow 和 JAX 成长起来的，PyTorch 用户想用上就得费更多功夫。
旧的方式
到目前为止，在 TPU 上运行 PyTorch 意味着要经过 PyTorch/XLA（torch_xla）。它能用，但有些怪癖。张量是惰性的（说的是张量，不是你），所以你的代码并不总是在你以为的时候运行。小小的改动就可能触发计算图重新编译，调试起来也很难推理。
为了让生活更轻松，我们构建了 Optimum TPU。它让 HF Mirror 模型在 TPU 上获得推理和训练支持。既然 TorchTPU 带来了原生支持，用 Transformers 和 Diffusers 运行 HF Mirror 模型就不再需要 Optimum TPU 了。
TorchTPU 登场 🚀
TorchTPU为 PyTorch 新增了 TPU 后端。你照常写 PyTorch 代码。它自带为 TPU 调优的 Pallas 内核、对分布式训练的支持，以及一个torch.compile集成，供你想要更快速度时使用。在 TPU 虚拟机上安装它只需一条命令：
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"。就这样。
from transformers import AutoModelForCausalLM, AutoTokenizer
model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="tpu", # 将直接在 TPU 上加载模型
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer(“AI 的未来是”, 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")。对于部署，编译模式只需增加一次compile调用。你的流水线和代码保持不变，也无需维护 TPU 分支。
import torch
from diffusers import StableDiffusion3Pipeline
pipe = StableDiffusion3Pipeline.from_pretrained(
"stabilityai/stable-diffusion-3.5-medium", dtype=torch.bfloat16
)
pipe.to("tpu") # eager 模式：这就是全部改动
# 可选：编译模式
pipe.transformer.compile(backend="tpu", fullgraph=True, dynamic=False)
image = pipe("一只柯基宇航员在披萨月球上插旗").images[0]
首次调用会进行编译，之后每次使用相同图像尺寸的调用都会复用该计算图。
模型太大，单芯片放不下？传入一个TensorParallelConfig给from_pretrained，transformer 在加载时就会分片到多个芯片上，因此完整模型无需全部装在一处。许多热门流水线开箱即支持，不支持的也只需一次tp_plan即可。Diffusers TPU 指南提供了完整配方。
以下是在 Cloud TPU v6e-4 上运行的一些示例：
TP=4（每个模型拆分到四个芯片上）· bf16 · batch 1
| MiniMax-H3（视频 + 音频） | Qwen-Image-2.1 | |
|---|---|---|
 |
 |
|
| "图片中那只毛茸茸的 HF Mirror 吉祥物，一张圆圆的黄色笑脸，脸颊红润，两只小手张开，站在一艘古希腊三列桨战船的船头 […] 海浪拍击声、木材吱呀声、风中索具声、远处海鸥声，以及低沉史诗般的管弦乐渐强。"· 1344×768，5.2 s，49 步 · 编译模式 188 s，eager 模式 262 s | "一只戴着巫师帽的水豚，在烛光下读书，油画风格"· 1024×1024，40 步 · 编译模式 3.5 s，eager 模式 8.0 s | "一块霓虹灯招牌，上面写着 'QWEN IMAGE 2.1'，雨夜，湿漉漉的路面上有倒影"· 2048×2048，40 步 |
内核
TheKernels Hub提供了可直接使用的 TPU 内核。用 Kernels 库加载一个，为你的模型提速。
TPU 内核用Pallas编写，让你能对硬件进行底层控制。你也可以引入来自Tokamax库的内核。下面是一个来自 Tokamax 的 ragged-dot 内核，接入到一个混合专家模型中：
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 仪表盘）[https://hf-mirror.com/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"。
我们迫不及待想看到你构建的东西 🤗。如果你遇到问题，请提交 issue 并告诉我们。