Hacker News 中文摘要

RSS订阅

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

原文链接 | HN讨论 | 2026-10-21 03:18:02

TorchTPU 为谷歌（Google）的张量处理单元（Tensor Processing Units，TPUs）提供 PyTorch 后端，使用户能够在 TPU 设备上运行 PyTorch 模型与算子。项目包含面向多种 PyTorch 操作的自定义 ATen 内核、用于优化 XLA 计算执行的编译缓存，并支持数据并行与张量并行的分布式训练；同时提供 Llama、通义千问（Qwen）、ResNet、DLRM、minGPT、UNet 等模型示例，并集成 torch.compile() 以实现图模式执行，另有调试与基准测试工具。

正文

描述

本项目为 Google 的 Tensor Processing Units（TPU）提供 PyTorch 后端，让用户能在 TPU 设备上运行 PyTorch 模型和操作。它包含用于多种 PyTorch 操作的自定义 ATen 内核、用于优化执行的编译缓存，以及分布式训练和模型推理的示例。

功能

  • 面向 TPU 的 PyTorch ATen 内核实现。
  • 用于优化 XLA 计算执行的编译缓存。
  • 支持分布式训练（数据并行、张量并行）。
  • 多种模型的示例（Llama、Qwen、ResNet、DLRM、minGPT、UNet）。
  • 与 torch.compile() 集成以支持图模式执行。
  • 用于调试和基准测试的实用工具。

用户指南与交互式 Notebook

用户指南见 https://torchtpu.readthedocs.io/。它是交互式 Notebook 的静态版本，Notebook 位于 docs 文件夹。这些 Notebook 基于 Marimo（类似 Jupyter），可用 marimo edit file.py 运行，然后在浏览器中访问。需要在 TPU VM 上运行，使用 v6e GCE 实例最简单。

安装

上面链接的用户指南中也提供了相同说明。TorchTPU 的 Python 版本支持策略与 PyTorch 的 Release Compatibility Matrix 一致，目前支持 Python 3.11+。详见 Python Version Support Policy。

首先创建 Python 虚拟环境：运行 python -m venv .venv，然后 source .venv/bin/activate；如果使用 uv，则运行 uv venv .venv，再 source .venv/bin/activate。

通过 pip 安装最新的公开发布版（PyPI 上的 0.2.0；如果使用 uv，请在所有 pip 命令前加上 uv）：pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

也可以通过源码构建 wheel：bazel build --config=wheelcommon //ci/wheel:torchtpu_wheel。

所有 OSS Bazel 构建都会通过 .bazelrc 自动默认使用优化模式（-c opt）。未优化构建（-c fastbuild、-c dbg）会降低性能，不用于 wheel 构建。如果你在开发或调试 C++ 代码，并明确想要未优化的 wheel，请传入 --//:allowunoptimizedwheel=True：bazel build -c dbg --config=wheelcommon //ci/wheel:torchtpuwheel --//:allowunoptimized_wheel=True。

必须保留 --config=wheelcommon，缺少它会在构建时失败。该配置会设置 --//:wheelbuild=True，使共享的 XLA/MLIR 后端经过 torchversion 重置转换，从而让 pywrap 将其整合为单个 libxlabase.so；若不加该标志，每个按版本划分的 common 都会附带自己的后端副本，生成的 wheel 在 import torch 时会因静态注册重复而中止。

wheelcommon 还会开启 Remote Build Execution（RBE）远程缓存和远程执行器。如果没有 RBE 凭据，可追加 --config=norbe 来去掉这些标志：bazel build --config=wheelcommon //ci/wheel:torchtpuwheel --config=norbe。注意顺序：--config=norbe 必须放在 --config=wheelcommon 之后，因为对于两者都设置的远程缓存和远程执行器标志，后出现的配置会生效。

安装 wheel：进入 <仓库路径>，运行 python -m venv .venv 并 source .venv/bin/activate，然后执行 pip install ../bazel-bin/ci/wheel/*.whl --extra-index-url https://download.pytorch.org/whl/cpu。这里的 --extra-index-url 提供 CPU 版 torch，而不是默认的 CUDA 构建。如果不加该参数，会安装 CUDA 版 torch；可以按上面方式传参，或单独安装 CPU 构建。

快速开始

以下是一些入门指引。

依赖

项目使用 Bazel 进行依赖管理。关键 Python 依赖：TBD。

Bazel 设置

要安装 Bazel，实际安装的是 Bazelisk——一个负责下载并安装 Bazel 本身的透明封装工具。常见安装方式有几种：如果已安装 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 进行验证：它会输出仓库信息且不报错。

VSCode clangd 配置

进行 C++ 代码导航时，建议使用 VSCode 配合 clangd 扩展。clangd 依赖由 Bazel action graph 生成的编译数据库，为编辑器提供代码补全、跳转定义等智能功能。

可通过界面或以下命令安装 VSCode 扩展：code --install-extension llvm-vs-code-extensions.vscode-clangd。同时卸载 Microsoft C++ 扩展以避免干扰：code --uninstall-extension ms-vscode.cpptools。生成初始 compilecommands.json：tools/setupclangd.py。

首次运行可能需要超过 30 分钟来完成完整构建，但重跑是增量的且很快。如果最近运行过 bazel build，生成的文件已存在，可传 --no-build 跳过等待（但在全新工作区中跳过构建可能导致 clangd 无法解析生成文件）。更多选项见脚本文件级注释。

也可以通过 nox 调用该脚本（-- 之后的参数会被转发）：nox -s refreshcompilecommands -- --no-build，或 nox -s refreshcompilecommands -- //src/torch_tpu/ops/...

贡献

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

星标历史

许可证

本项目采用 Apache 2.0 许可证。详见 LICENSE 文件。

评论

社区讨论以追问和产业观察为主，未出现强烈站队；分歧主轴是 TPU 及相关软件栈开放后，护城河会落在硬件、编译器还是云服务上。讨论者主要追问 TPU 能否走出 Google Cloud、硬件“铲子”是否泛滥以及对前沿实验室和云厂商的冲击，少数人乐观认为 XLA 编译器栈会普惠所有硬件，个别人猜测这是对 DeepSeek 开源推理引擎的回应。

认为 TPU 已通过云和实体合作走出 Google Cloud，硬件“铲子”将大量供给

  • 有人先问 Google Cloud 之外是否真有 TPU，显示外界对 TPU 实际可得性并不清楚（少数人，提问）
  • 回应称 Google 既通过云服务卖 TPU 算力，也向部分合作方实体出售 TPU 硬件；因此相关开放顺理成章，硬件“铲子”会变得充足，Nebius 股价已从炒作高点跌近一半（少数人）

猜测此举是 Google 对 DeepSeek 开源推理引擎的回应

  • 一位非 AI 工程师提问，这是否是 Google 针对 DeepSeek 开源推理引擎的举措（少数人） > 这是不是 Google 对 DeepSeek 开源推理引擎的回应？（@yipinwong）

认为模型和推理层都不再是护城河后，硬件正成为新的护城河

  • 若模型不是护城河、推理层也不是护城河，竞争壁垒可能已转移到硬件；这引发对前沿实验室未来处境的疑问（少数人） > 如果模型和推理层都不是护城河，那护城河是不是现在变成硬件了？（@nelaggy）

认为 XLA 是比 Inductor 更好的编译器基础设施，并将加速 TensorFlow 退场

  • XLA 本身不错，且可能比 PyTorch 的 Inductor 更适合作为编译器基础设施，从而提升所有硬件上的性能；这也像是给 TensorFlow 钉上最后一颗钉子（少数人） > XLA 可能比 Inductor 更好，会提升所有硬件性能，也给 TensorFlow 钉上最后一颗钉子。（@fancyfredbot）

有价值的补充

  • npalli 指出亚马逊在 AWS Trainium 上也有同类做法，并给出 aws-neuron/torch-neuronx 作为 PyTorch 集成的参考。（@npalli）
  • pcmoritz 提到自己在做面向 Tenstorrent 加速器的开源 libtpu 实现，项目叫 libtt，基于 Tenstorrent 软件栈，已能运行大部分 torchtpu，并展示把 PyTorch 下沉到分块式加速器所涉及的工作。（@pcmoritz）