Google 釋出 TorchTPU：PyTorch 原生 TPU 後端

說明

本專案提供 PyTorch 後端以支援 Google 的 Tensor Processing Units (TPUs)，讓使用者能在 TPU 裝置上執行 PyTorch 模型與操作。它包含針對各種 PyTorch 操作的自訂 ATen 核心、最佳化執行的編譯快取，以及分散式訓練與模型推論的範例。

功能

PyTorch ATen 核心實作於 TPU。
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 的 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 安裝

安裝最新公開版本 (0.2.0 on PyPI)（對於 uv，在所有 pip 指令前加上 uv）：

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；若未使用此旗標，則每個版本共用其自身的後端複製，導致 wheel 在 import torch 上因重複靜態註冊而中止。

wheel_common 同時啟用 Remote Build Execution (RBE) 遠端快取與遠端執行器。若您沒有 RBE 憑證，請加上 --config=no_rbe 以移除這些旗標：

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

順序很重要：--config=no_rbe 必須在 --config=wheel_common 之後，因為後者設定的 remote-cache 與 remote-executor 旗標會覆蓋前者。

使用以下方式安裝 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 安裝，將安裝 CUDA 版的 torch；如上所示傳遞 --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 能夠為編輯器提供各種智慧功能，包括程式碼補完、跳轉至定義、等。

透過 UI 或以下指令安裝 VSCode 擴充功能：

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

產生初始 compile_commands.json：
```
tools/setup_clangd.py
```

第一次執行可能需要一段時間（> 30 分鐘）才能完成完整建構，但之後的執行將是增量且費用低。若您最近已執行過 bazel build，已產生的檔案存在，您可以傳遞 --no-build 以跳過等待大型建構（請注意，在全新工作空間中跳過建構可能會阻止 clangd 解析已產生的檔案）。更多選項請參閱 script 檔案級別註解。

另外，您也可以透過 nox 呼叫指令碼（-- 之後的引數會被轉發）：

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

貢獻

我們歡迎對 PyTorch on TPU 專案的貢獻！請參閱我們的 CONTRIBUTING.md，瞭解如何貢獻的指引，包括簽署我們的 Contributor License Agreement (CLA) 與程式碼審查流程。

Star 歷史

授權

本專案採用 Apache 2.0 授權。詳情請參閱 LICENSE 檔案。