Google 已开源 TorchTPU，这是一个专门设计的后端，旨在让 PyTorch 模型和操作直接在 Google 的 Tensor Processing Units (TPUs) 上运行。虽然 PyTorch 在历史上一直依赖各种变通方法来与 TPU 硬件交互，但 TorchTPU 通过包含专门为 TPU 架构实现的自定义 ATen 内核（PyTorch 操作的基础构建块）提供了一种更完整的集成体验。这种原生支持旨在简化研究人员和工程师的开发工作流，让他们能够在无需非原生集成带来的摩擦的情况下，利用 Google 的专用 AI 加速器。
这张图片标识了 TorchTPU 项目，该项目为 PyTorch 模型在 Google Tensor Processing Units 上运行提供原生后端。
该后端引入了多个技术特性来优化性能和可扩展性。它包含一个编译缓存，用于加速 XLA（加速线性代数）计算的执行，而 XLA 是 TPU 处理机器学习工作负载的核心。此外，TorchTPU 支持分布式训练技术，包括数据并行（Data Parallelism）和张量并行（Tensor Parallelism），这些技术对于训练 Llama 或 Qwen 等大规模模型至关重要。该项目还与 `torch.compile()` 集成，允许用户受益于图模式执行优化，这可以在训练和推理过程中显著提高吞吐量并降低延迟。
星数历史图表说明了开发者对这一全新的 TPU 原生 PyTorch 后端的兴趣日益浓厚。
尽管转向原生支持，该项目仍强调了实现最佳使用的特定硬件和软件要求。例如，文档指出，虽然它支持各种模型，但交互式笔记本和用户指南最好在 TPU VM（例如 v6e GCE 实例）上运行。此外，由于该项目依赖于 XLA/MLIR 后端，这意味着虽然它提供了更无缝的 PyTorch 体验，但它仍运行在更广泛的 XLA 生态系统中。这一进展标志着 PyTorch 生态系统中对 TPU 的第一方支持的转变，有可能缩小 JAX 的硬件中心化优化与 PyTorch 更广泛的软件采用之间的差距。