Hugging Face 推出 TorchTPU：实现 PyTorch 在 TPU 上的原生支持
此前，若在 TPU 上运行 PyTorch，开发者通常需要将代码移植到 torch_xla，或转而学习 JAX。随着 TorchTPU 的发布，这一情况正在改变。TorchTPU 作为 PyTorch 的原生后端，将 TPU 视为标准设备，从而开启了 Hugging Face 生态系统对 TPU 的原生支持。Transformers、Diffusers 和 Kernels 库均可在 TPU 上原生运行。用户只需从 Hub 选择模型，将其发送至 "tpu" 即可使用。
什么是 TPU？
TPU（Tensor Processing Unit，张量处理单元）是 Google 自研的 AI 加速器。这类芯片专为神经网络中常见的矩阵运算设计，Google 已研发此类硬件约十年，并用于训练和服务其内部模型，包括 Gemini。
TPU 通过 Google Cloud 提供，规模可从单芯片扩展至数千芯片。在需要大规模计算的场景下，无论是模型训练还是大规模服务，TPU 均能发挥显著优势。
长期以来，软件生态是 TPU 应用的瓶颈。现有的 TPU 软件栈主要基于 TensorFlow 和 JAX 构建。
传统方案的局限
过去，在 TPU 上运行 PyTorch 依赖 PyTorch/XLA (torch_xla)。尽管该方案可行，但存在诸多限制。代码执行往往不符合预期，微小的改动可能触发图重新编译，且调试过程复杂难懂。
为简化流程，Hugging Face 曾推出 Optimum TPU，为 Transformers 和 Diffusers 模型提供推断和训练支持。如今，TorchTPU 带来了原生支持，使得在 TPU 上运行 Hugging Face 模型不再必须依赖 Optimum TPU。
TorchTPU 核心特性
TorchTPU 将 TPU 集成为 PyTorch 的新后端。用户无需修改代码逻辑，编写的是标准的 PyTorch 代码。它支持分布式训练，并集成 torch.compile 以提升性能。在 TPU 虚拟机上安装仅需一条命令：
pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu
额外的索引 URL 确保获取 CPU 版本的 PyTorch 构建包。若不指定此参数，Pip 可能会拉取体积更大且不需要的 CUDA 版本。
主流库在 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",
)
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")。对于生产部署，可添加编译调用以优化性能。管道和代码结构保持不变，无需维护针对 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]
首次调用会触发编译，后续相同图像尺寸的调用将复用编译后的图表。
若模型过大无法容纳于单个芯片，可向 from_pretrained 传递 TensorParallelConfig。变压器将在加载时跨芯片分割权重，确保完整模型无需驻留于单一位置。许多流行管道已支持此功能，其他管道可通过配置 tp_plan 实现。详细配置参考 Diffusers TPU 指南。
以下示例数据来自云 TPU v6e-4：
- 配置：TP=4（模型分布于四个芯片）、bf16 精度、Batch Size=1
| MiniMax-H3 (视频 + 音频) | Qwen Image 2.1 | |
|---|---|---|
| 提示词：“图片中的毛毛拥抱面孔吉祥物...海浪声...史诗管弦乐” · 分辨率 1344×768, 时长 5.2 秒, 49 步 · 编译耗时 188 秒, Eager 模式耗时 262 秒 | 提示词：“一个戴着巫师帽子的巴拉，在光下阅读一本书，油画” · 分辨率 1024×1024, 40 步 · 编译耗时 3.5 秒, Eager 模式耗时 8.0 秒 | 提示词：“一个读到 'QWEN IMAGE 2.1' 的商店标志，雨夜，湿路上的反射” · 分辨率 2048×2048, 40 步 |
Kernels
Kernels Hub 提供了现成的 TPU 内核。通过 Kernels 库加载内核可加速模型推理。这允许用户对硬件进行底层控制，既可使用 Pallas 编写自定义 TPU 内核，也可从 Tokamax 库中获取现成内核。
以下是将 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 Size 为 64 时，Qwen3-30B-A3B 的解码速度比默认实现快 21.5 倍，Gemma 4 26B-A4B 快 12.8 倍。（测试环境：TPU v6e-8，对比 Hub Kernels 与默认实现）。
测试与稳定性
团队通过运行库测试套件持续改进 Transformers 中的 TPU 支持。目前正在推进 TPU 夜间 CI 的建设，结果将在 Transformers CI Dashboard 中展示。
目前，涉及 TPU 使用的测试通过率已达 96.5%。发现的每一个故障都被转化为小型复现脚本，并作为 Issue 提交至 TorchTPU 仓库。截至目前，已报告 34 个问题，其中 19 个已解决。此外，TPU 测试还发现了 Transformers 本身的一些错误（如 StaticCache 和 Qwen3 张量并行计划中的问题），这些已在上游修复。
开始使用
对于 PyTorch 用户而言，TPU 不再是边缘任务。安装 TorchTPU，从 Hub 获取模型，并将其发送至 "tpu" 即可。
- GitHub: TorchTPU
- Hugging Face Hub: TPU Kernels
- 文档：Cloud TPU Documentation
