TorchTPU：PyTorch 原生 TPU 后端（官方博客中文摘要）

TorchTPU 是 Google 与 Meta 之间的开源合作项目，为 Google TPU 引入一个原生的、设计规范的 PyTorch 后端。它支持开箱即用的模型执行与 TPU 迁移，并通过 eager 执行优化、torch.compile 集成和原生硬件加速，提供具有竞争力的推理与训练性能。

背景

Google TPU 长期在 Google 内外部被成功用于大规模机器学习模型的训练与部署，但历史上主要与 TensorFlow 和 JAX 生态关联。2020 年，Google 推出了首个面向 TPU 的编译版 PyTorch，即基于 Lazy Tensors 的 PyTorch/XLA。TorchTPU 是 2025 年中在 Google 启动、与 Meta PyTorch 团队合作的新项目，把 eager 执行作为首要目标：现有 PyTorch 工作负载只需极小的代码改动（例如把 .to("cuda") 改为 .to("tpu")）即可开箱运行。

当前进展

TorchTPU 仍处于积极开发中，已被部分大型 Google Cloud 客户用于生产环境的推理与训练负载，并以开源项目形式发布。目前已支持 eager 与编译两种模式，支持约 450 个 Aten 算子（涵盖主流深度学习架构所需核心 ATen 算子的前向与反向计算），并支持通过 JAX/Pallas 和 Helion DSL 编写自定义内核。

数值精度

对 TPU 而言，原生的矩阵乘法格式是 bfloat16（BF16）。TorchTPU 默认采用 medium 精度模式：乘法使用 BF16，累加保持在 FP32。对绝大多数深度学习工作负载，项目团队建议使用这一默认精度。

来源与许可

本文基于 PyTorch 官方博客文章改写：原文《Introducing TorchTPU》由 Ebba Simpson 发表在 pytorch.org，按 CC BY 3.0 许可发布（阅读原文）。本文为中文摘要改写，未新增原文之外的事实与数字；技术细节以原文为准。

核心要点

• TorchTPU 是 Google 与 Meta 之间的开源合作项目
• 现有 PyTorch 工作负载只需极小的代码改动（例如把 .to("cuda") 改为 .to("tpu")）即可开箱运行
• 支持约 450 个 Aten 算子
• TorchTPU 默认采用 medium 精度模式

常见问题

TorchTPU 是什么？

为 Google TPU 引入一个原生的、设计规范的 PyTorch 后端

TorchTPU 现在可以用于生产环境吗？

已被部分大型 Google Cloud 客户用于生产环境的推理与训练负载