Google开源TorchTPU：PyTorch原生TPU后端
Google open-sources TorchTPU, a native PyTorch back end for TPUs
PyTorch用户上TPU的官方原生方案落地，解决算力适配痛点，值得关注部署细节。
Description
描述
This project provides a PyTorch backend for Google's Tensor Processing Units (TPUs), enabling users to run PyTorch models and operations on TPU devices. It includes custom ATen kernels for various PyTorch operations, a compilation cache for optimizing execution, and examples for distributed training and model inference.
本项目为 Google 的张量处理单元（TPU）提供 PyTorch 后端，使用户能够在 TPU 设备上运行 PyTorch 模型和运算。它包含针对各种 PyTorch 运算的自定义 ATen 内核、用于优化执行的编译缓存，以及分布式训练和模型推理的示例。
Features
功能特性
- PyTorch ATen kernel implementation for TPU.
- Compilation cache for optimized execution of XLA computations.
- Support for distributed training (Data Parallelism, Tensor Parallelism).
- Examples for various models (Llama, Qwen, ResNet, DLRM, minGPT, UNet).
- Integration with torch.compile() for graph mode execution.
- Utilities for debugging and benchmarking.
- TPU 的 PyTorch ATen 内核实现。
- 用于优化 XLA 计算执行的编译缓存。
- 支持分布式训练（数据并行、张量并行）。
- 多种模型的示例（Llama、Qwen、ResNet、DLRM、minGPT、UNet）。
- 与 torch.compile() 集成以支持图模式执行。
- 用于调试和基准测试的工具。
User Guide and Interactive Notebooks
用户指南与交互式笔记本
A user guide is available at https://torchtpu.readthedocs.io/. This is a static version of the interactive notebooks. They are available in the docs folder.
用户指南位于 https://torchtpu.readthedocs.io/。这是交互式笔记本的静态版本，它们也存在于 docs 文件夹中。
The notebooks are Marimo (similar to Jupyter). They can be run with marimo edit file.py and then you can access them in your browser. You will need to be running this from a TPU VM, a v6e GCE instance is easiest.
这些笔记本基于 Marimo（类似于 Jupyter）。你可以使用 marimo edit file.py 运行它们，然后在浏览器中访问。你需要在 TPU VM 上运行此操作，v6e GCE 实例最为简便。
Installation
安装
The same instructions are available in the user guide linked above.
上述链接的用户指南中也提供了相同的说明。
TorchTPU's Python version support policy matches PyTorch's Release Compatibility Matrix. Currently, Python 3.11+ is supported. See the Python Version Support Policy for more details.
TorchTPU 的 Python 版本支持策略与 PyTorch 的版本兼容性矩阵保持一致。目前支持 Python 3.11+。有关更多详细信息，请参阅 Python 版本支持策略。
First create a Python virtual environment:
首先创建一个 Python 虚拟环境：
python -m venv .venv
source .venv/bin/activateor, with uv
或者，使用 uv
uv venv .venv
source .venv/bin/activateInstall via pip
通过 pip 安装
Install the latest public release (0.2.0 on PyPI) (for uv prepend uv to all pip commands):
安装最新的公开版本（PyPI 上的 0.2.0）（对于 uv，在所有 pip 命令前添加 uv）：
pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpuBuild from source
从源代码构建
Wheels can be built via:
可以通过以下方式构建 Wheels：
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheelAll OSS Bazel builds default to optimized mode (-c opt) automatically via .bazelrc. Unoptimized builds (-c fastbuild, -c dbg) can degrade performance and are not used for wheel builds. If you are developing or debugging C++ code and explicitly want an unoptimized wheel, pass --//:allow_unoptimized_wheel=True:
所有开源 Bazel 构建默认通过 .bazelrc 自动设置为优化模式（-c opt）。未优化的构建（-c fastbuild、-c dbg）可能会降低性能，且不会用于 wheel 构建。如果你正在开发或调试 C++ 代码并明确需要未优化的 wheel，请传递 --//:allow_unoptimized_wheel=True：
bazel build -c dbg --config=wheel_common //ci/wheel:torch_tpu_wheel \
--//:allow_unoptimized_wheel=True--config=wheel_common is required, and the build fails at build time without it. The config sets --//:wheel_build=True, which routes the shared XLA/MLIR backend through the torch_version-reset transition so that pywrap factors it into a single libxla_base.so; built without the flag, every per-version common ships its own copy of the backend and the resulting wheel aborts on import torch with duplicate static registrations.
必须使用 --config=wheel_common，否则构建将在构建时失败。该配置设置了 --//:wheel_build=True，这将通过 torch_version-reset 转换将共享的 XLA/MLIR 后端路由到 pywrap 中，使其将其整合为单个 libxla_base.so；如果不使用该标志构建，每个版本的公共部分都会自带其自己的后端副本，导致在 import torch 时因重复的静态注册而中止。
wheel_common also turns on the Remote Build Execution (RBE) remote cache and remote executor. If you do not have RBE credentials, append --config=no_rbe to strip those flags back out:
wheel_common 还会启用远程构建执行（RBE）的远程缓存和远程执行器。如果你没有 RBE 凭据，请追加 --config=no_rbe 以移除这些标志：
bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbeThe order matters here: --config=no_rbe has to come after --config=wheel_common, because the later config is the one that wins for the remote-cache and remote-executor flags they both set.
这里的顺序很重要：--config=no_rbe 必须放在 --config=wheel_common 之后，因为对于它们共同设置的 remote-cache 和 remote-executor 标志，后面的配置会生效。
Install wheel via:
通过以下方式安装 wheel：
cd <path_to_repo>
python -m venv .venv; source .venv/bin/activate
# The extra-index-url provides the CPU version of torch rather than the default CUDA build.
pip install ../bazel-bin/ci/wheel/*.whl --extra-index-url https://download.pytorch.org/whl/cpuNote that installing without --extra-index-url https://download.pytorch.org/whl/cpu will install the CUDA version of torch; pass --extra-index-url as shown above or install the CPU build separately.
注意，如果不使用 --extra-index-url https://download.pytorch.org/whl/cpu 进行安装，将会安装 torch 的 CUDA 版本；请按照上述方式传递 --extra-index-url，或者单独安装 CPU 构建版本。
Getting Started
入门指南
Some pointers to get you started:
一些入门提示：
- PyTorch: Official PyTorch documentation
- Tutorials and examples: easy to understand PyTorch code using TorchTPU
- Public APIs: the TorchTPU APIs that you can use
- PyTorch：官方 PyTorch 文档
- 教程和示例：使用 TorchTPU 编写的易于理解的 PyTorch 代码
- 公共 API：你可以使用的 TorchTPU API
Dependencies
依赖项
The project uses Bazel for dependency management. Key Python dependencies are:
该项目使用 Bazel 进行依赖管理。主要的 Python 依赖项包括：
TBD
待补充
Bazel setup
Bazel 设置
This project uses Bazel for building and dependency management.
本项目使用 Bazel 进行构建和依赖管理。
To install Bazel, we actually install Bazelisk, which is a transparent wrapper that handles downloading and installing Bazel itself.
为了安装 Bazel，我们实际上安装的是 Bazelisk，它是一个透明的包装器，负责处理 Bazel 本身的下载和安装。
There are several ways to install Bazelisk. Here are some common ways:
有几种安装 Bazelisk 的方法。以下是一些常见的方式：
- If you have Go installed:
- 如果你已安装 Go：
go install github.com/bazelbuild/bazelisk@latest- Download the binary directly:
- 直接下载二进制文件：
wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64See the Bazelisk README for more install instructions
请参阅 Bazelisk README 以获取更多安装说明
However you install it, add it's location to PATH.
无论你如何安装它，请将其位置添加到 PATH 中。
To verify the install, run bazel info in the repo's root directory:
要验证安装是否成功，请在仓库根目录下运行 bazel info：
bazel infoIt should print information about the repo without errors.
它应该打印出关于仓库的信息且没有错误。
VSCode clangd Setup
VSCode clangd 设置
For C++ code navigation, it is recommended to use VSCode with clangd extension. Relying on a compilation database generated from Bazel action graph, clangd is able to enrich the editor with various smart features, including code completions, go-to definitions, etc.
对于 C++ 代码导航，建议使用 VSCode 配合 clangd 扩展。clangd 依赖于从 Bazel 操作图生成的编译数据库，能够为编辑器提供丰富的智能功能，包括代码补全、跳转到定义等。
- Install VSCode extensions via UI or the following command:
- 通过 UI 或以下命令安装 VSCode 扩展：
code --install-extension llvm-vs-code-extensions.vscode-clangd
# Uninstall Microsoft C++ extensions to avoid interference.
code --uninstall-extension ms-vscode.cpptools- Generate the initial compile_commands.json:
- 生成初始的 compile_commands.json：
tools/setup_clangd.pyIt may take a while (> 30 minutes) for the first run to make a full build, but reruns will be incremental and cheap. If you have run bazel build recently, generated files already exist so you can pass --no-build to skip waiting for heavy builds (note that on a fresh workspace, skipping the build may prevent clangd from resolving generated files). See more options in the script file-level comments.
首次运行可能需要一段时间（> 30 分钟）来完成完整构建，但后续运行将是增量且快速的。如果你最近运行过 bazel build，生成的文件已经存在，因此你可以传递 --no-build 参数以跳过等待重型构建的过程（注意：在全新的工作区中，跳过构建可能会导致 clangd 无法解析生成的文件）。有关更多选项，请参阅脚本文件级别的注释。
Alternatively, you can invoke the script via nox (arguments after -- are forwarded):
或者，你可以通过 nox 调用该脚本（-- 之后的参数将被转发）：
nox -s refresh_compile_commands -- --no-build
nox -s refresh_compile_commands -- //src/torch_tpu/ops/...Contributing
贡献指南
We welcome contributions to the PyTorch on TPU project! Please see our CONTRIBUTING.md for guidelines on how to contribute, including signing our Contributor License Agreement (CLA) and code review processes.
我们欢迎对 PyTorch on TPU 项目的贡献！请参阅我们的 CONTRIBUTING.md，了解如何做出贡献的指导方针，包括签署我们的贡献者许可协议（CLA）以及代码审查流程。
Star History
星标历史
License
许可证
This project is licensed under the Apache 2.0 License. See the LICENSE file for details.
本项目采用 Apache 2.0 许可证授权。详情请参阅 LICENSE 文件。
更进一步：量化金融体系
看懂新闻只是起点——沿量化金融路径，把它变成能交付的工程能力