Hugging Face 生态接入 TorchTPU，TPU 成为 PyTorch 原生设备
HF Mirror ❤️ TorchTPU: TPUs, Natively
Hugging Face 宣布与 TorchTPU 合作，让 PyTorch 把 TPU 当作普通设备使用，Transformers、Diffusers 和 Kernels 可在 TPU 上原生运行。
原文给出了 TPU 原生接入后的安装方式、代码改动量和实测算例，读者可以据此判断现有 PyTorch 工作流的迁移成本。
如果你曾经尝试过在 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 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]
首次调用会进行编译，之后每次用相同图像尺寸调用都会复用计算图。
模型太大，单颗芯片放不下？把一个 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 秒，49 步 · 编译耗时 188 秒，即时执行 262 秒 | "一只戴着巫师帽的水豚，在烛光下读书，油画风格" · 1024×1024，40 步 · 编译耗时 3.5 秒，即时执行 8.0 秒 | "一块写着‘QWEN IMAGE 2.1’的霓虹灯招牌，雨夜，湿漉漉的路面上倒映着灯光" · 2048×2048，40 步 |
Kernels
Kernels Hub 提供了开箱即用的 TPU kernel。使用 Kernels 库加载一个，即可为你的模型提速。
TPU kernel 使用 Pallas 编写，让你能够对硬件进行底层控制。你也可以引入 Tokamax 库中的 kernel。下面是一个来自 Tokamax 的 ragged-dot kernel，已接入一个 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 kernel 与默认实现的对比。
我们先替你测了 🧪
我们通过运行库的测试套件，不断改进 transformers 对 TPU 的支持。我们正在努力添加基于 TPU 的 nightly CI，并将结果发布到 (transformers dashboard)[https://hf-mirror.com/spaces/transformers-community/transformers-ci-dashboard] 中。
Transformers CI dashboard，TPU 结果与其他加速器并列展示。仍在进行中！
并非所有测试都已通过，我们也接受这一点。在 TPU 上运行的测试中，96.5% 通过。我们发现的每个失败都会变成一个小的复现脚本和 TorchTPU 仓库上的一个 issue。到目前为止，我们已经报告了 34 个 issue，其中 19 个已经修复。在 TPU 上运行还发现了 Transformers 自身的一些 bug，例如在 StaticCache 和 Qwen3 张量并行计划中的问题。我们已在上游修复了这些问题。
如果 TPU 上出现问题，我们希望比你更早发现。
开始使用
TPU 对 PyTorch 用户来说不再是一个支线任务。安装 TorchTPU，从 Hub 获取一个模型，然后把它发送到 "tpu"。
我们迫不及待想看到你构建的作品 🤗。如果遇到问题，请提交 issue 告诉我们。
来源：Hugging Face Blog · hf-mirror.com