Hugging Face 发布 TorchTPU，让 PyTorch 原生支持 TPU
Hugging Face ❤️ TorchTPU: TPUs, Natively
Hugging Face 发布 TorchTPU，一个原生 PyTorch 后端，让 TPU 成为 PyTorch 的普通设备，Hugging Face 生态的 Transformers、Diffusers 和 Kernels 均可原生运行于 TPU。
PyTorch用户过去需借助torch_xla或JAX才能用TPU，TorchTPU原生后端让Hugging Face模型指定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 上获取，从单芯片到由数千个芯片组成的集群。当你需要大量算力用于训练或大规模部署时，它们表现出色。
难点一直在于软件。TPU 软件栈是围绕 TensorFlow 和 JAX 发展起来的，PyTorch 用户需要付出更多努力才能接入。
旧方式
直到现在，在 TPU 上运行 PyTorch 意味着要通过 PyTorch/XLA（torch_xla）。它可行，但有其怪癖。张量是惰性的（张量是惰性的，不是你），所以你的代码并不总是在你以为它运行时运行。微小的更改可能会触发图重新编译，且调试难以推理。
为了简化操作，我们构建了 Optimum TPU。它为 Hugging Face 模型提供了在 TPU 上的推理和训练支持。现在 TorchTPU 带来了原生支持，运行 Hugging Face 模型的 Transformers 和 Diffusers 不再需要 Optimum TPU。
TorchTPU 登场 🚀
TorchTPU 将 TPU 作为新后端添加到 PyTorch。你编写普通的 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 mode）运行。你也可以使用 torch.compile 以获得更好的性能，或在分布式环境中运行。
图像？视频？Diffusers 在 TPU 上绘制它们 🎨
TPU 现在是 Diffusers 中的一等公民设备。对于原型设计或调试，急切模式只需要 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 · 批次 1
| MiniMax-H3（视频 + 音频） | Qwen-Image-2.1 | |
|---|---|---|
 |
 |
|
| “图片中那只毛茸茸的 Hugging Face 吉祥物，一张圆圆的黄色笑脸，脸颊红润，伸出两只小小的手掌，站在一艘古希腊三列桨战船（trireme）的船头 […] 海浪拍击声、木材吱呀声、索具中的风声、远处的海鸥声，以及低沉而宏大的管弦乐渐强。” · 1344×768，5.2 秒，49 步 · 编译耗时 188 秒，急切模式耗时 262 秒 | “一只戴着巫师帽的水豚，在烛光下读书，油画风格” · 1024×1024，40 步 · 编译耗时 3.5 秒，急切模式耗时 8.0 秒 | “一个写着 'QWEN IMAGE 2.1' 的霓虹灯招牌，雨夜，湿漉漉的路面上的倒影” · 2048×2048，40 步 |
内核
The 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 的张量并行方案中的问题。我们已在上游修复了这些问题。
如果 TPU 上出现问题，我们希望在你发现之前就能找到。
开始使用
对于 PyTorch 用户来说，TPU 不再是一个支线任务。安装 TorchTPU，从 Hub 获取一个模型，然后将其发送到 "tpu"。
我们迫不及待想看看你会构建什么 🤗。如果你遇到问题，请提交一个 issue 并告诉我们。
来源：Hugging Face Blog · huggingface.co