Hugging Face 支持 TorchTPU：Transformers 与 Diffusers 原生运行在 TPU 上
Hugging Face ❤️ TorchTPU: TPUs, Natively
Hugging Face 宣布 TorchTPU 作为 PyTorch 原生后端接入其生态，Transformers、Diffusers 和 Kernels 可直接在 TPU 上运行，不再需要 Optimum TPU。
原文给出 TorchTPU 接入 Transformers 与 Diffusers 的具体用法和实测数据，可据此判断 PyTorch 在 TPU 上的迁移成本。
如果你曾经尝试在 TPU 上运行 PyTorch，你很可能把代码移植到了 torch_xla，或者最终不得不学习 JAX。借助 TorchTPU 这一原生 PyTorch 后端，TPU 现在只是另一个设备。它为 Hugging Face 生态带来了对 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。它让 Hugging Face 模型在 TPU 上获得推理和训练支持。既然 TorchTPU 带来了原生支持，用 Transformers 和 Diffusers 运行 Hugging Face 模型就不再需要 Optimum TPU 了。
TorchTPU 登场 🚀
TorchTPU 为 PyTorch 新增了 TPU 后端。你照常编写普通的 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 · 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 kernel。用 Kernels 库加载一个，就能让你的模型提速。
TPU kernel 用 Pallas 编写，让你能对硬件进行底层控制。你也可以从 Tokamax 库引入 kernel。下面是一个来自 Tokamax 的 ragged-dot kernel，被接入到一个混合专家模型中：
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 kernel 对比默认实现。
我们替你测试，你就不用测了 🧪
我们通过运行库的测试套件，持续改进 transformers 对 TPU 的支持。我们正在努力为 TPU 添加每日 CI，并将结果发布在 (transformers dashboard)[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard] 中。
Transformers CI dashboard，TPU 结果与其他加速器并列展示。仍在进行中！
目前并非全部通过，我们对此也能接受。在 TPU 上运行的测试中，96.5% 通过。我们发现的每一个失败都会变成一个小的复现脚本和 TorchTPU 仓库上的一个 issue。到目前为止，我们已报告 34 个 issue，其中 19 个已修复。在 TPU 上运行还发现了 Transformers 自身的一些 bug，比如 StaticCache 和 Qwen3 张量并行计划中的问题。我们在上游修复了这些。
如果 TPU 上出了问题，我们希望比你更早发现。
开始使用
对 PyTorch 用户来说，TPU 不再是支线任务。安装 TorchTPU，从 Hub 上取一个模型，然后把它发送到 "tpu"。
我们迫不及待想看到你构建的东西 🤗。如果遇到问题，请开一个 issue 告诉我们。
来源：Hugging Face Blog · huggingface.co