跳到正文
原文
Hacker News· ilreb·· 3 小时前精选AI 评分68

Google 开源原生 PyTorch 后端 TorchTPU

Google open-sources TorchTPU, a native PyTorch back end for TPUs

AI 导读

Google 开源 TorchTPU，这是一套面向 Google TPU 的原生 PyTorch 后端，可让用户在 TPU 设备上运行 PyTorch 模型和操作。

推荐理由

仓库集中交代了后端能力、分布式训练支持、示例模型以及安装构建路径，便于读者评估 PyTorch 在 TPU 上的集成范围与落地成本。

正文 · AI 翻译

TorchTPU Logo

描述

本项目为 Google 的 Tensor Processing Units (TPUs) 提供 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

通过 pip 安装

安装最新的公开发行版（PyPI 上的 0.2.0）（对于 uv，将 uv 添加在所有 pip 命令之前）：

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

从源代码构建

可通过以下方式构建 wheel：

bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel

所有 OSS Bazel 构建都会自动采用优化模式（-c opt），具体通过 .bazelrc 实现。未优化构建（-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，使共享的 XLA/MLIR 后端通过 torch_version-reset 转换，从而让 pywrap 将其整合为单个 libxla_base.so；如果在没有此标志的情况下构建，每个按版本构建的 common 都会自带一份后端副本，最终生成的 wheel 会因静态注册重复而在 import torch 处中止。

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
# 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/cpu

请注意，安装时如果不使用 --extra-index-url https://download.pytorch.org/whl/cpu，将安装 torch 的 CUDA 版本；如上所示传入 --extra-index-url，或者单独安装 CPU 构建版本。

入门

以下是一些入门指引：

依赖项

该项目使用 Bazel 进行依赖管理。主要的 Python 依赖项包括：

待定

Bazel 设置

本项目使用 Bazel 进行构建和依赖管理。

要安装 Bazel，我们实际上安装的是 Bazelisk，它是一个透明封装器，负责处理 Bazel 本身的下载和安装。

安装 Bazelisk 有多种方式。以下是一些常见方式：

  • 如果你已安装 Go：

    go install github.com/bazelbuild/bazelisk@latest

  • 直接下载二进制文件：

    wget -O ~/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.27.0/bazelisk-linux-amd64

有关更多安装说明，请参阅 Bazelisk README

无论采用哪种安装方式，请将其位置添加到 PATH。

要验证安装，请在仓库根目录运行 bazel info：

bazel info

它应该会输出有关该仓库的信息，且不会出现错误。

VSCode clangd 设置

对于 C++ 代码导航，建议使用 VSCode，并搭配 clangd 扩展。 依靠从 Bazel 操作图生成的 编译数据库，clangd 能够 为编辑器提供各种智能功能，包括代码补全、跳转定义、 等等。

  1. 通过 UI 或以下命令安装 VSCode 扩展：

    code --install-extension llvm-vs-code-extensions.vscode-clangd
# Uninstall Microsoft C++ extensions to avoid interference.
code --uninstall-extension ms-vscode.cpptools

  2. 生成初始 compile_commands.json：

    tools/setup_clangd.py

首次运行进行完整构建可能需要一段时间（> 30 分钟），但 后续运行将采用增量构建，成本也很低。如果你最近运行过 bazel build， 生成的文件已经存在，因此可以传入 --no-build，跳过等待 耗时的构建（请注意，在全新工作区中，跳过构建可能会导致 clangd 无法解析生成的文件）。有关更多选项，请参阅 脚本 文件级注释。

或者，你可以通过 nox 调用脚本（-- 之后的参数会被转发）：

nox -s refresh_compile_commands -- --no-build
nox -s refresh_compile_commands -- //src/torch_tpu/ops/...

贡献

欢迎为 PyTorch on TPU 项目贡献力量！有关贡献指南，包括签署我们的贡献者许可协议（CLA）和代码审查流程，请参阅我们的 CONTRIBUTING.md。

Star History

许可证

本项目采用 Apache 2.0 许可证授权。详情请参阅 LICENSE 文件。

来源：Hacker News · github.com

#开源/仓库#部署/工程#Google