Hugging Face 联合 TorchTPU 让 PyTorch 原生支持 TPU
Hugging Face ❤️ TorchTPU: TPUs, Natively
Hugging Face 宣布通过 PyTorch 原生后端 TorchTPU，Transformers、Diffusers 和 Kernels 无需 torch_xla 或 Optimum TPU 即可原生运行在 TPU 上，模型只需 device_map="tpu" 或 pipe.to("tpu")。
官方给出了 Transformers 与 Diffusers 原生跑 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 虚拟机上安装它只需一条命令：
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 调用。你的 pipeline 和代码保持不变，也不用维护 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 会在加载时被分片到多个芯片上，因此完整模型永远不必装进单处。许多流行的 pipeline 开箱即支持，不支持的也只需一个 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 核心。用 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 的张量并行计划中的问题。我们已在上游修复了这些 bug。
如果在 TPU 上出了问题，我们希望在你之前发现它。
开始使用
对 PyTorch 用户来说，TPU 不再是支线任务。安装 TorchTPU，从 Hub 上找一个模型，然后把它送到 "tpu" 上。
我们迫不及待想看到你构建的东西 🤗。如果遇到问题，请开一个 issue 并告诉我们。
来源： Hugging Face Blog · huggingface.co