Hugging Face ❤️ TorchTPU
来自 Hugging Face · 期号 2026-10-20
摘要
在 TPU 上跑 PyTorch 过去要走 PyTorch/XLA，张量是惰性的，调试起来很绕。 TorchTPU 把 TPU 变成了 PyTorch 的原生后端。装一个包，模型加载时加一句 `device_map="tpu"`，就能在 TPU 上跑 Transformers 和 Diffusers。 Diffusers 那边更简单：`pipe.to("tpu")` 就能用 eager 模式跑文生图，编译模式再加一次 `compile` 调用。模型太大放不进单芯片时，传一个 `TensorParallelConfig` 就能自动分片。 Kernels Hub 提供了现成的 TPU kernel，用 Pallas 写的。在 Qwen3-30B-A3B 上 batch size 64 时解码快了 21.5 倍。 Hugging Face 正在跑 CI 测试覆盖，目前 TPU 上 96.5% 的测试通过，已向 TorchTPU 仓库报了 34 个 issue，其中 19 个已修复。
深读
导读 TorchTPU 将 Google TPU 变成了 PyTorch 的原生后端，终结了 PyTorch 用户必须通过 PyTorch/XLA 或转学 JAX 才能使用 TPU 的局面。Hugging Face 的 Transformers 和 Diffusers 库现已原生支持 TPU，只需一行 device_map="tpu" 即可运行。读完本文你将了解如何零改动地将现有模型迁移到 TPU，以及 Kernels Hub 上现成的 Pallas 内核能带来多大的性能提升。 关键信息 - TorchTPU 通过 pip install torch-tpu --extra-index-url https://download.pytorch.org/whl/cpu 一条命令安装，无需 CUDA 构建。 - Transformers v5 中加载模型到 TPU 只需在 from_pretrained 中传入 device_map="tpu"，默认以 eager 模式运行。 - Diffusers 中 eager 模式仅需 pipe.to("tpu")，编译模式再加一次 pipe.transformer.compile(backend="tpu") 调用；模型过大时传入 TensorParallelConfig 即可在加载时自动分片。 - Kernels Hub 提供用 Pallas 编写的 TPU 内核；在 Qwen3-30B-A3B 上 batch size 64 时解码吞吐量提升 21.5 倍，在 Gemma 4 26B-A4B 上提升 12.8 倍。 - Hugging Face 正在为 Transformers 添加 TPU 上的 nightly CI，目前 96.5% 的测试通过，已向 TorchTPU 仓库报告 34 个 issue，其中 19 个已修复。 逐节讲解 背景：TPU 是什么，过去怎么用 TPU（Tensor Processing Unit）是 Google 自研的 AI 加速器，专为神经网络最耗时的矩阵运算设计，已迭代约十年，用于训练和部署 Gemini 等模型。TPU 通过 Google Cloud 提供，从单芯片到数千芯片的 pod 均可使用，适合大规模训练和推理。其软件栈过去围绕 TensorFlow 和 JAX 构建，PyTorch 用户只能通过 PyTorch/XLA（torch_xla）接入，但张量是惰性的（代码不一定会按预期立即执行），小改动可能触发图重编译，调试困难。Hugging Face 此前为此构建了 Optimum TPU 库，但 TorchTPU 实现原生支持后已不再需要。 TorchTPU：TPU 成为 PyTorch 原生后端 TorchTPU 将 TPU 作为新后端加入 PyTorch，支持标准 Py…