Today for AI
返回全部资讯
Google 开源 TorchTPU 实现 PyTorch 原生 TPU 推理Hacker News AIT2·开源与评测

Google 开源 TorchTPU 提供 PyTorch TPU 原生后端

原标题：Google open-sources TorchTPU, a native PyTorch back end for TPUs
78AI 研判分

Google 正式开源 TorchTPU，为 PyTorch 框架提供原生的 Tensor Processing Units (TPU) 后端支持。该工具实现了自定义 ATen 内核与编译缓存，支持分布式训练及 torch.compile() 集成，旨在降低开发者在 TPU 上运行 Llama、Qwen 等主流模型的技术门槛。

报道正文

本文目录9 个章节
<img src="docs/images/torchtpu_logo.avif" alt="TorchTPU Logo" style="width: 60%;">

描述

本项目为 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 虚拟环境：

BASH
python -m venv .venv
source .venv/bin/activate

或者，使用 uv：

BASH
uv venv .venv
source .venv/bin/activate

通过 pip 安装

安装最新的公开版本（PyPI 上的 0.2.0）（如果使用 uv，请在所有 pip 命令前加上 uv）：

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

从源码构建

可以通过以下命令构建 Wheel 包：

BASH
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：

BASH
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 重置转换（transition），以便 pywrap 将其整合为单个 libxla_base.so；若不带此标志进行构建，每个版本特定的公共组件都会包含一份后端的独立副本，导致生成的 wheel 包在 import torch 时因重复的静态注册而中止。

wheel_common 还会启用远程构建执行（RBE）的远程缓存和远程执行器。如果你没有 RBE 凭证，请追加 --config=no_rbe 以移除这些标志：

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

这里的顺序很重要：--config=no_rbe 必须位于 --config=wheel_common 之后，因为对于两者都设置的远程缓存和远程执行器标志，后出现的配置具有优先级。

通过以下方式安装 wheel 包：

BASH
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

该命令应输出有关仓库的信息且无错误。

（注：更多技术实现细节与完整 API 文档请查阅原项目 README）