Hugging Face 原生支持 TPU：TorchTPU 让 PyTorch 代码无需改动直接运行

2026-10-21 12:46 0 阅读

从 torch_xla 到原生支持

如果你试过在 TPU 上跑 PyTorch，大概率经历过两种选择：把代码移植到 torch_xla，或者干脆去学 JAX。现在有了 TorchTPU，TPU 变成了 PyTorch 的又一个普通设备。Hugging Face 的整个生态——Transformers、Diffusers、Kernels——都能原生运行。从 Hub 拉个模型，指定 device_map="tpu"，完事。

TPU 是什么，为什么之前不好用

TPU（Tensor Processing Unit）是 Google 自研的 AI 加速器，专门干一件事：神经网络里占大头的那部分矩阵运算。Google 已经迭代了约十年，自家模型包括 Gemini 的训练和推理都跑在上面。Cloud TPU 从单芯片到数千芯片的 Pod 都有，大规模训练或高吞吐推理是它的主场。

但软件栈一直是门槛。TPU 生态围绕 TensorFlow 和 JAX 长大，PyTorch 用户得额外费不少劲。以前唯一的办法是走 PyTorch/XLA（torch_xla），能跑，但体验别扭：张量是"惰性"的，代码执行时机跟直觉不符；小改动可能触发整图重编译；调试起来也很难定位问题。

Hugging Face 之前做了 Optimum TPU 来缓解这些痛点，给模型提供推理和训练支持。现在 TorchTPU 带来原生后端，Optimum TPU 不再是必需品。

TorchTPU 怎么用

TorchTPU 把 TPU 注册为 PyTorch 的新后端。你写普通的 PyTorch 代码就行。它内置了针对 TPU 调优的 Pallas kernel，支持分布式训练，也接入了 torch.compile 做加速。

安装只需要一条命令，在 TPU VM 上执行：

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

注意后面那个 extra index，它指定了 CPU 版本的 torch。不加的话 pip 会拉体积大得多的 CUDA 版本，在 TPU 上完全没必要。

Transformers：一行代码迁移

Transformers v5 的设计哲学是成为模型定义的单一可信来源，让它们能在任何地方跑。"任何地方"现在正式包括 TPU。

加载模型时传 device_map="tpu" 即可：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Qwen/Qwen3.5-9B"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="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 现在是 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")
# 可选：编译模式
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（视频+音频）	1344×768，49 步，5.2 秒	编译后 188 秒，eager 262 秒
Qwen-Image-2.1	1024×1024，40 步	编译后 3.5 秒，eager 8.0 秒
Qwen-Image-2.1	2048×2048，40 步	—

Kernels：即插即用的加速

Kernels Hub 已经有现成的 TPU kernel。用 Kernels 库加载，就能给模型提速。

TPU kernel 用 Pallas 编写，可以精细控制硬件。也可以从 Tokamax 库引入 kernel。下面是把 Tokamax 的 ragged-dot kernel 插进 MoE 模型的例子：

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",
)

batch 越大收益越明显。batch size 64 时，Qwen3-30B-A3B 的解码吞吐提升 21.5 倍，Gemma 4 26B-A4B 提升 12.8 倍（TPU v6e-8 上对比默认实现）。

测试进展：96.5% 通过率，问题在修

Hugging Face 团队在持续跑 Transformers 的完整测试套件来打磨 TPU 支持，也在搭建 TPU nightly CI，结果会放到 Transformers CI dashboard。

目前还没全绿，但团队接受这个状态。TPU 上跑的测试通过率 96.5%，每个失败案例都会变成最小复现脚本和 TorchTPU repo 上的 issue。目前已经报了 34 个，19 个已修复。TPU 测试还帮 Transformers 本身抓到了几个 bug，比如 StaticCache 和 Qwen3 张量并行计划的问题，都已 upstream 修复。

目标很明确：问题在开发者遇到之前就解决掉。