Hugging Face 推出 TorchTPU，为 PyTorch 原生支持 TPU
Hugging Face ❤️ TorchTPU: TPUs, Natively
Hugging Face 推出 TorchTPU，为 PyTorch 增加原生 TPU 后端，Transformers、Diffusers 和 Kernels 可在 TPU 上运行。
TorchTPU 将 TPU 接入原生 PyTorch 工作流，并覆盖 Transformers、Diffusers 与 Kernels；文中还提供测试通过率和性能数据，呈现了支持范围与当前进展。
如果你曾尝试在 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 已经设计 TPU 约十年，并用它们训练和部署自家的模型，包括 Gemini。
TPU 可通过 Google Cloud 使用，规模从单个芯片到数千个芯片组成的 pod 不等。当你需要大量算力进行训练或大规模部署时，TPU 尤其出色。
一直以来，难点都在软件上。TPU 技术栈围绕 TensorFlow 和 JAX 构建，PyTorch 用户要想使用 TPU，就得多费些功夫。
以往的方式
直到现在，要在 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]
首次调用会进行编译，此后每次使用相同图像尺寸调用，都会复用计算图。
模型太大，单个芯片放不下？将 TensorParallelConfig 传给 from_pretrained，加载时就会把 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 s，49 步 · 编译模式 188 s，急切模式 262 s | “戴着巫师帽的水豚在烛光下读书，油画” · 1024×1024，40 步 · 编译模式 3.5 s，急切模式 8.0 s | “霓虹商店招牌上写着 ‘QWEN IMAGE 2.1’，雨夜，湿漉漉的人行道上倒映着灯光” · 2048×2048，40 步 |
内核
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 上添加夜间 CI，并在 (transformers 仪表板)[https://huggingface.co/spaces/transformers-community/transformers-ci-dashboard] 上公布结果。
Transformers CI 仪表板，展示 TPU 与其他加速器的测试结果。仍在开发中！
目前测试结果还不是全绿，但我们对此并不介意。在 TPU 上运行的测试中，96.5% 通过。我们发现的每个失败都会变成一个简短的复现脚本，并在 TorchTPU 仓库中提交一个问题。目前我们报告了 34 个问题，其中 19 个已修复。在 TPU 上运行测试还发现了 Transformers 本身的一些 bug，例如 StaticCache 和 Qwen3 张量并行方案中的问题。我们已在上游修复了这些问题。
如果 TPU 上出现问题，我们希望能抢在你之前发现。
开始使用
对于 PyTorch 用户来说，TPU 不再是支线任务。安装 TorchTPU，从 Hub 获取模型，然后将其发送到 "tpu"。
我们迫不及待地想看看你会构建出什么 🤗。如果遇到问题，请提交 issue 告诉我们。
来源：Hugging Face Blog · huggingface.co