Hugging Face TorchTPU 原生集成

PyTorch 中的原生 TPU 支持

TorchTPU 实现了对 Google 张量处理单元 (TPU) 的原生 PyTorch 支持，消除了对 torch_xla 或 JAX 等复杂变通方案的需求。通过将 TPU 集成为标准设备后端，开发者现在可以使用熟悉的 PyTorch 语法直接在 TPU 硬件上运行 Hugging Face Transformers 和 Diffusers 模型。只需将设备设置为 "tpu"，模型即可在即时 (eager) 模式下执行，并可选择使用 torch.compile 以获得更高的性能。

超越传统工作流

从历史上看，在 TPU 上运行 PyTorch 需要 torch_xla，这带来了诸如惰性张量评估、复杂的调试以及频繁的图重新编译等挑战。TorchTPU 的引入取代了对之前 Optimum TPU 库的需求，提供了更无缝的体验。安装只需通过一条命令即可完成：

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

这种方法使用 torch 的 CPU 构建版本以避免不必要的 CUDA 依赖，从而为 TPU 硬件优化环境。

与 Transformers 和 Diffusers 的集成

Transformers v5 现在将 TPU 视为一等公民。用户可以通过将 device_map="tpu" 传递给 from_pretrained，直接将模型加载到 TPU 硬件上。对于 Diffusers，流水线通过标准的 .to("tpu") 方法支持 TPU。

对于超过单芯片内存容量的大规模模型，用户可以通过传递 TensorParallelConfig 来实现张量并行。这会在加载过程中将模型分片到多个芯片上，从而能够在 TPU pod 上执行大规模架构。

高性能内核

用户可以通过利用来自 Hugging Face Kernels Hub 的专用内核进一步优化 TPU 性能。这些内核使用 Pallas 编写，提供底层硬件控制。例如，使用来自 Tokamax 的 ragged-dot 内核处理混合专家 (MoE) 模型，已证明能显著提高吞吐量。在批处理大小为 64 时，Qwen3-30B-A3B 和 Gemma 4 26B-A4B 等模型的解码速度与默认实现相比，性能分别提升了 21.5 倍和 12.8 倍。

持续测试与可靠性

为了确保稳定性，Hugging Face 正在将 TPU 支持集成到其持续集成 (CI) 流水线中。目前，96.5% 的测试库套件可以在 TPU 硬件上通过。团队正在积极跟踪故障，向 TorchTPU 存储库报告问题，并将修复程序上游合并到 Transformers 库中，以提高整体兼容性。这项持续的努力旨在使 TPU 部署像传统的基于 GPU 的工作流一样可靠。