Google 开源 TorchTPU，这是 TPU 的原生 PyTorch 后端

该项目为 Google 张量处理单元 (TPU) 提供 PyTorch 后端，使用户能够在 TPU 设备上运行 PyTorch 模型和操作。它包括用于各种 PyTorch 操作的自定义 ATen 内核、用于优化执行的编译缓存以及分布式训练和模型推理的示例。 TPU 的 PyTorch ATen 内核实现。用于优化 XLA 计算执行的编译缓存。支持分布式训练（数据并行、张量并行）。各种模型的示例（Llama、Qwen、ResNet、DLRM、minGPT、UNet）。与 torch.compile() 集成以进行图形模式执行。用于调试和基准测试的实用程序。用户指南和交互式笔记本 用户指南可从 https://torchtpu.readthedocs.io/ 获取。这是交互式笔记本的静态版本。它们可以在 docs 文件夹中找到。笔记本是 Marimo（类似于 Jupyter）。它们可以使用 marimo edit file.py 运行，然后您可以在浏览器中访问它们。您需要从 TPU VM 运行它，v6e GCE 实例是最简单的。上面链接的用户指南中提供了相同的说明。 TorchTPU 的 Python 版本支持政策与 PyTorch 的版本兼容性矩阵相匹配。目前支持Python 3.11+。有关更多详细信息，请参阅 Python 版本支持政策。首先创建一个 Python 虚拟环境： python -m venv .venv source .venv/bin/activate 或使用 uv uv venv .venv source .venv/bin/activate 安装最新的公开版本（ PyPI 上的 0.2.0 ）（对于 uv 将 uv 添加到所有 pip 命令）： pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu Wheels 可以通过以下方式构建： bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel 所有 OSS 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 是必需的，如果没有它，构建将在构建时失败。配置集 --//:wheel_build=True ，它通过 torch_version -reset 转换路由共享 XLA/MLIR 后端，以便 pywrap 将其分解为单个 libxla_base.so ；在没有该标志的情况下构建，每个版本的公共版本都会附带自己的后端副本，并且生成的轮子在导入火炬时会因重复的静态注册而中止。 wheel_common 还打开远程构建执行（RBE）远程缓存和远程执行器。如果您没有 RBE 凭据，请附加 --config=no_rbe 以删除这些标志： bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel --config=no_rbe 这里的顺序很重要： --config=no_rbe 必须在 --config=wheel_common 之后，因为后面的配置是赢得它们同时设置的远程缓存和远程执行器标志的配置。通过以下方式安装wheel： cd < path_to_repo > python -m venv .venv ; source .venv/bin/activate # extra-index-url 提供 Torch 的 CPU 版本，而不是默认的 CUDA 版本。 pip install ../bazel-bin/ci/wheel/ * .whl --extra-index-url https://download.pytorch.org/whl/cpu 请注意，不带 --extra-index-url https://download.pytorch.org/whl/cpu 安装将安装 CUDA 版本的 torch ；如上所示传递 --extra-index-url 或单独安装 CPU 版本。一些入门指南：该项目使用 Bazel 进行依赖管理。主要的 Python 依赖项有： TBD 该项目使用 Bazel 进行构建和依赖项管理。为了安装 Bazel，我们实际上安装了 Bazelisk，它是一个透明的包装器，负责下载和安装 Bazel 本身。有多种方法可以安装 Bazelisk。以下是一些常见的方法： 请参阅 Bazelisk 自述文件以获取更多安装说明 无论您如何安装它，请将其位置添加到 PATH 中。要验证安装，请在存储库的根目录中运行 bazel info：它应该打印有关存储库的信息，没有错误。对于 C++ 代码导航，建议使用带有 clangd 扩展的 VSCode。依靠 Bazel 操作图生成的编译数据库，clangd 能够为编辑器提供各种智能功能，包括代码补全、转到定义等。通过 UI 或以下命令安装 VSCode 扩展： code --install-extension llvm-vs-code-extensions.vscode-clangd # 卸载 Microsoft C++ 扩展以避免干扰。 code --uninstall-extension ms-vscode.cpptools 生成初始的compile_commands.json：第一次运行可能需要一段时间（> 30 分钟）才能完成完整构建，但重新运行将是增量且便宜的。如果您最近运行了 bazel build，生成的文件已经存在，因此您可以传递 --no-build 来跳过等待繁重的构建（请注意，在新的工作空间上，跳过构建可能会阻止 clangd 重新构建）