Hugging Face 生态通过 TorchTPU 原生支持 TPU

如果你曾经尝试在 TPU 上运行 PyTorch，你大概要么把代码移植到了 torch_xla，要么最终不得不去学 JAX。有了 TorchTPU 这个原生 PyTorch 后端，TPU 现在只是另一个设备而已。它让 Hugging Face 生态实现了对 TPU 的原生支持。Transformers、Diffusers 和 Kernels 都能在 TPU 上原生运行。从 Hub 上选一个模型，把它发到 "tpu"，就搞定了。

等等，TPU 是什么？🤔

TPU 是 Tensor Processing Unit（张量处理单元）的缩写。TPU 是谷歌自家的 AI 加速器。它们是为一项任务而生的芯片：神经网络大部分时间都在做的矩阵运算。谷歌设计它们已经大约十年了，并用它们来训练和部署自家的模型，包括 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 将 TPU 作为新后端加入 PyTorch。你编写常规的 PyTorch 代码。它自带为 TPU 调优的 Pallas 内核、对分布式训练的支持，以及一个 torch.compile 集成，供你想要更高速度时使用。在 TPU VM 上安装它只需一条命令：

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", # 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")。对于部署，编译模式只需额外一次 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 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（每个模型拆分到四颗芯片上）· bf16 · 批大小 1}

MiniMax-H3（视频 + 音频）	Qwen-Image-2.1
“图片中那个毛茸茸的 Hugging Face 吉祥物，一张圆圆的黄色笑脸，脸颊泛红，两只小手张开，站在一艘古希腊三列桨座战船的船首 […] 海浪拍击声、木材吱嘎作响、风吹动索具、远处海鸥鸣叫，以及低沉史诗般的管弦乐渐强。” · 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 步

Kernels

Kernels Hub 提供了可直接使用的 TPU 内核。用 Kernels 库加载一个，给你的模型提速。

TPU 内核用 Pallas 编写，它让你对硬件进行底层控制。你也可以从 Tokamax 库引入内核。下面是一个来自 Tokamax 的 ragged-dot 内核，被接入到一个混合专家（Mixture-of-Experts）模型中：

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 上出了问题，我们希望比你先发现。

开始使用

TPU 对 PyTorch 用户来说不再是小众玩法。安装 TorchTPU，从 Hub 上拿一个模型，然后把它发送到 "tpu"。

• GitHub 上的 TorchTPU
• Hub 上的 TPU 内核
• Cloud TPU 文档

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