TPU 代表张量处理单元。 TPU 是谷歌自己的人工智能加速器。它们是为一项工作而构建的芯片：神经网络将大部分时间花在矩阵数学上。谷歌已经设计它们大约十年了，并使用它们来训练和服务自己的模型，包括 Gemini。

TPU 可在 Google Cloud 上使用，从单个芯片到数千个芯片的 Pod。当您需要大量计算、进行培训或大规模服务时，它们就会发挥作用。

问题一直是软件。 TPU 堆栈是围绕 TensorFlow 和 JAX 发展起来的，PyTorch 用户必须更加努力才能进入。

到目前为止，在 TPU 上运行 PyTorch 意味着要通过 PyTorch/XLA (torch_xla)。它有效，但有其怪癖。张量是惰性的（张量，而不是你），所以你的代码并不总是在你认为它运行时运行。小的变化可能会触发图形重新编译，并且调试很难推理。

为了让生活更轻松，我们打造了 Optimum TPU。它为 Hugging Face 模型提供了 TPU 上的推理和训练支持。现在，TorchTPU 带来了原生支持，不再需要 Optimum TPU 来运行带有 Transformers 和 Diffusers 的 Hugging Face 模型。

TorchTPU 将 TPU 添加为 PyTorch 的新后端。您编写常规 PyTorch 代码。它配备了针对 TPU 进行调整的 Pallas 内核、对分布式训练的支持以及在您需要更高速度时的 torch.compile 集成。在 TPU VM 上安装它需要一个命令：

额外的索引为您提供了 torch 的 CPU 版本。如果没有它，pip 会拉动更大的 CUDA 构建，而您在 TPU 上不需要它。

Transformers v5 围绕一个想法构建：成为模型定义的真实来源，以便它们可以在任何地方运行。 “Anywhere”现在包括 TPU。要在 TPU 上加载模型，请传递 device_map="tpu"。就是这样。

该模型默认以 Eager 模式运行。您还可以使用 torch.compile 以获得更好的性能，或在分布式设置中运行它。

TPU现在是扩散器中的一流器件。对于原型设计或调试，Eager 模式只需要 pipeline.to("tpu")。对于部署，编译模式添加了一个编译调用。您的管道和代码保持不变，并且无需维护 TPU 分支。

第一次调用会编译，此后具有相同图像大小的每次调用都会重用该图。

对于一颗芯片来说模型太大？将 TensorParallelConfig 传递给 from_pretrained，变压器在加载时会跨芯片分片，因此完整模型永远不必放在一个位置。许多流行的管道都支持它开箱即用，而那些不支持的管道则只需 tp_plan 即可。扩散器 TPU 指南有完整的配方。

以下是在 Cloud TPU v6e-4 上运行的一些示例：

TP=4（每个模型分为四个芯片）·bf16·批次 1

Kernels Hub 具有现成的 TPU 内核。使用内核库加载一个模型并提高模型速度。

TPU 内核是用 Pallas 编写的，它使您可以对硬件进行低级控制。您还可以从 Tokamax 库引入内核。这是一个来自 Tokamax 的参差不齐的内核，插入到专家混合模型中：

批次越大，胜利就越大。在批量大小为 64 时，Qwen3-30B-A3B 上的解码速度提高了 21.5 倍，Gemma 4 26B-A4B 上的解码速度提高了 12.8 倍。