Hugging Face 通过 TorchTPU 为 Transformers 与 Diffusers 提供原生 TPU 支持

如果你曾尝试在 TPU 上运行 PyTorch，你大概把代码移植到了 torch_xla，或者最终不得不去学 JAX。有了原生 PyTorch 后端 TorchTPU，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 把 TPU 作为新后端加入 PyTorch。你写的还是普通的 PyTorch 代码。它自带为 TPU 调优的 Pallas kernels、分布式训练支持，以及在你想要更快速度时可用的 torch.compile 集成。在 TPU VM 上安装它只需一条命令：

pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu

这个额外的 index 会给你 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]

第一次调用会进行编译，之后每次使用相同图像尺寸的调用都会复用该计算图。

模型大到一块芯片放不下？把 TensorParallelConfig 传给 from_pretrained，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 s, 49 steps · 188 s 编译, 262 s eager	"一只戴着巫师帽的水豚，在烛光下读书，油画" · 1024×1024, 40 steps · 3.5 s 编译, 8.0 s eager	"一块霓虹店招，上面写着 'QWEN IMAGE 2.1'，雨夜，湿路面上的倒影" · 2048×2048, 40 steps

内核

Kernels 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 添加 nightly CI，并在 (transformers 仪表板)[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 上出了问题，我们希望比你先发现。

开始使用

对 PyTorch 用户来说，TPU 不再是支线任务。安装 TorchTPU，从 Hub 取一个模型，然后把它送到 "tpu"。

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

我们迫不及待想看看你会做出什么 🤗。如果遇到问题，开一个 issue 告诉我们。