Google 开源 TorchTPU：在 TPU 上原生运行 PyTorch

torch_tpu 是面向 Google TPU 的 PyTorch 后端，让 PyTorch 模型和算子直接在 TPU 设备上运行。项目实现了 TPU 的 ATen 算子，支持 torch.compile() 图模式、数据并行和张量并行，并提供 Llama、Qwen、ResNet 等示例模型。

要点

• 已在 PyPI 发布 0.2.0 版本，支持 Python 3.11 及以上
• 带有 XLA 计算的编译缓存，以及调试和基准测试工具
• Apache 2.0 许可；README 未标注项目成熟度
• 前一天 AWS 刚为 Trainium 开源了原生 PyTorch 后端 TorchNeuron，云厂商自研芯片都在补齐 PyTorch 原生支持

来源

### [Google 开源 TorchTPU：在 TPU 上原生运行 PyTorch]({{URL}})

torch_tpu 是面向 Google TPU 的 PyTorch 后端，让 PyTorch 模型和算子直接在 TPU 设备上运行。项目实现了 TPU 的 ATen 算子，支持 torch.compile() 图模式、数据并行和张量并行，并提供 Llama、Qwen、ResNet 等示例模型。

- 已在 PyPI 发布 0.2.0 版本，支持 Python 3.11 及以上
- 带有 XLA 计算的编译缓存，以及调试和基准测试工具
- Apache 2.0 许可；README 未标注项目成熟度
- 前一天 AWS 刚为 Trainium 开源了原生 PyTorch 后端 TorchNeuron，云厂商自研芯片都在补齐 PyTorch 原生支持

**为什么值得关注**：过去在 TPU 上用 PyTorch 主要依赖 PyTorch/XLA，原生后端有望降低迁移成本；有 TPU 资源的团队可以评估现有训练代码的改动量。

来源：[google-pytorch/torch_tpu](https://github.com/google-pytorch/torch_tpu)；[Google open-sources TorchTPU, a native PyTorch back end for TPUs (HN)](https://news.ycombinator.com/item?id=50169053)