google-pytorch/torch_tpu

这是什么？

核心功能

• 01为 TPU 提供自定义 ATen 内核实现。
• 02编译缓存，复用已编译的 XLA 计算。
• 03支持数据并行和张量并行的分布式训练。
• 04可与 torch.compile() 集成，进行图模式执行。
• 05提供 Llama、Qwen、ResNet、UNet 等模型示例。
• 06包含调试和基准测试工具。

适用场景

• 将现有 PyTorch 训练代码迁移到 Google Cloud TPU 上运行。
• 在多个 TPU 设备上进行大模型分布式训练。
• 在 TPU 硬件上部署 PyTorch 模型推理。
• 学习常见模型在 TPU 上的示例实现。

快速上手

1. 1

   创建 Python 虚拟环境

   python -m venv .venv
2. 2

   从 PyPI 安装最新公开版本，并使用 CPU 版 PyTorch 索引

   pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

使用前须知

• !面向 Google TPU 硬件，只有能访问 TPU（如 TPU VM）时才有用。
• !需要 Python 3.11 或更高版本。
• !从源码构建需要 Bazel；完整 wheel 构建还需要 Remote Build Execution 凭据，否则需使用 no_rbe 参数。
• !元数据中的许可证显示为 NOASSERTION，但 README 声明为 Apache 2.0。

标签

常见问题

torch_tpu 是什么？

让 PyTorch 模型和算子运行在 Google TPU 上的 PyTorch 后端。 TorchTPU 为 Google TPU 提供 PyTorch 后端，使现有 PyTorch 代码可以在 TPU 设备上执行。项目包含针对多种 PyTorch 算子的自定义 ATen 内核，并提供编译缓存，用于复用已编译的 XLA 计算以加速重复执行。

torch_tpu 可以免费使用吗？用的什么许可证？

torch_tpu 以 Other 许可证开源。用于商业产品前，请先确认许可证条款。

torch_tpu 还在维护吗？

活跃维护中。torch_tpu 最近一次提交是在 2026年10月22日。

torch_tpu 适合用来做什么？

将现有 PyTorch 训练代码迁移到 Google Cloud TPU 上运行。 在多个 TPU 设备上进行大模型分布式训练。 在 TPU 硬件上部署 PyTorch 模型推理。 学习常见模型在 TPU 上的示例实现。

torch_tpu 怎么安装和运行？

创建 Python 虚拟环境: python -m venv .venv 从 PyPI 安装最新公开版本，并使用 CPU 版 PyTorch 索引: pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

在哪里获取 torch_tpu？

源代码和安装说明在 GitHub 上：https://github.com/google-pytorch/torch_tpu