Hugging Face发布TorchTPU：PyTorch原生支持谷歌TPU，无需再移植至JAX

Hugging Face于2026年10月20日发布TorchTPU，这是一个为PyTorch添加的原生张量处理单元（TPU）后端。这意味着开发者现在可以直接将模型发送到“tpu”设备上运行，而无需像过去那样必须将代码移植到torch_xla库，或者转而学习JAX框架。这一变化旨在降低谷歌云TPU在主流PyTorch生态系统中的使用门槛。

TPU是谷歌专为神经网络矩阵运算设计的AI加速器芯片，已用于训练其Gemini等大模型。长期以来，TPU的软件栈主要围绕TensorFlow和JAX构建，导致PyTorch用户在使用时面临诸多不便，例如张量的惰性执行使得代码运行时机难以预测，以及小改动触发图重编译带来的调试困难。TorchTPU通过提供常规PyTorch代码的直接支持、针对TPU优化的Pallas内核以及torch.compile集成，解决了这些痛点。

随着TorchTPU的发布，Hugging Face的核心库如Transformers v5和Diffusers现已原生支持TPU。开发者只需在安装TorchTPU后，加载模型时指定device_map="tpu"即可。对于图像和视频生成任务，Diffusers也允许通过简单的pipe.to("tpu")调用实现原型设计，或通过编译模式提升部署性能。此外，Kernels Hub提供了现成的TPU内核，可在混合专家模型（MoE）等场景中显著提升解码吞吐量。

尽管兼容性仍在持续改进中，Hugging Face报告称在TPU上运行的测试套件中有96.5%的用例通过，并已修复了包括StaticCache在内的多个上游Bug。这一进展使得TPU不再是PyTorch用户的“支线任务”，而是成为与GPU并列的一流计算选项，可能促使更多依赖谷歌云算力的企业和个人开发者转向TPU进行大规模模型训练和推理。