TorchTPU PyTorch TPU 指南：2026 年安裝與遷移方式

把 PyTorch 程式從 cuda 移植到 tpu，TorchTPU 讓整體改動相對精簡；不過 TPU 是否可用、套件選擇、編譯流程與張量形狀，才是能否順利執行與取得效能的關鍵。

TorchTPU 真正能做到什麼

TorchTPU 是 Google TPU 的原生 PyTorch 後端，不是代管模型 API，也不是 TPU VM 的替代方案。PyTorch 官方公告提到 eager execution、torch.compile、分散式原語，以及約 450 個 ATen 運算的支援；公開儲存庫則列出 ATen kernels、編譯快取、分散式範例與 torch.compile 整合。

遷移時的目標刻意維持熟悉的寫法：

# Before
model.to("cuda")

# First TPU port
model.to("tpu")

但這只能證明程式具備可攜性，不能保證效能。PyTorch 公告也指出，TPU 能否發揮峰值效能，仍取決於模型形狀、分片方式，以及運算是否適合這套硬體。

安裝前需要準備什麼

只有在 Google TPU 已完成配置的環境中，TorchTPU 才真正派得上用場。官方儲存庫建議在 TPU VM 上執行 Notebook，並將 v6e Google Compute Engine 執行個體列為最容易開始的環境。只有 CPU 的筆電可以安裝套件，卻無法驗證 TPU 是否真的執行。

公開 README 支援 Python 3.11 以上版本，並指定 torch_tpu==0.2.0。建議使用隔離環境，避免實驗沿用 CUDA 版本或其他無關相依套件。

安裝 `torch_tpu`，避免拉入 CUDA 套件

先建立虛擬環境，再搭配指定版本的 TorchTPU 套件安裝 CPU 版 PyTorch wheel：

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

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

README 特別提醒，省略額外套件索引可能會拉入 CUDA 版 PyTorch；對 TPU VM 來說，這並不是正確的相依套件。

第一次遷移建議直接使用已發布的 wheel。若要從原始碼建置，則需要 Bazel 或 Bazelisk，以及儲存庫文件所列的 --config=wheel_common，必要時還要加上 --config=no_rbe。

用最小化 PyTorch 程式確認後端

請在 TPU VM 上執行以下程式：

import torch

x = torch.randn(2, 2, device="tpu")
y = torch.randn(2, 2, device="tpu")
z = x @ y

print(z)
print(z.device)

如果程式成功輸出結果，且裝置顯示為 tpu，就代表張量配置與矩陣乘法已交由 TPU 後端處理。若裝置不可用，先檢查 VM 執行環境與套件安裝，不要急著修改模型程式。

TorchTPU 同時支援方便除錯的 eager execution，以及用於編譯執行的 torch.compile。PyTorch 公告指出，根據專案測試，DeferAndFuse 模式相較嚴格 eager mode 可能快 50% 以上；但這是該專案的測試結果，不應視為適用所有工作負載的通用基準。

把實際的 PyTorch 工作負載遷移過去

先完成最小幅度的修改，確認單一裝置流程可正常運作後，再加入 TPU 專用的分散式設定。

將裝置選擇改為 torch.device("tpu")；在支援的環境中，也可以使用 torch.accelerator.current_accelerator()。
確認模型與輸入張量都移到同一個裝置。
啟用編譯或分散式執行前，先跑一個短批次。
測量時維持批次大小與張量形狀穩定；形狀變更可能觸發編譯工作。
在正式投入較大規模遷移前，先測試具代表性的模型流程與分散式模式。

分散式 CUDA 程式通常會使用 NCCL；TorchTPU 發布的範例則改用 TPU process-group 後端：

import torch
import torch.distributed as dist

device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")
model = model.to(device)

實際的分散式啟動指令，會視 TPU 環境與工作負載而定。只改變裝置字串，並不會自動完成分片、程序配置或輸入管線設定。

第一次執行的結果與後續執行可能不同。TorchTPU 文件提到，編譯快取可以重複利用相同的計算圖，因此應將預熱時間與穩態吞吐量分開記錄。

TorchTPU 哪些地方已經成熟，哪些仍需留意

Google 已公開 TorchTPU 與 vLLM、SGLang、TorchTitan 及 Hugging Face 工具鏈的整合方式。其推論公告介紹了 vLLM 與 SGLang 的原生 TPU 路徑，包括 KV-cache 搬移與服務功能。Google 的 Cloud 公告則報告 TorchTitan 執行 Qwen 3 8B 時達到 52% MFU 的結果，以及其在正式環境中的使用情況。

不過，實際使用時的限制，沒有公告標題看起來那麼完整：

面向	來源已確認的內容	你仍應自行確認的事項
安裝	Python 3.11+、`torch_tpu==0.2.0`、CPU PyTorch 索引	你的 VM 映像與 wheel 是否能順利解析相依套件
可攜性	Eager mode、`torch.compile`、`tpu` 裝置、分散式支援	不支援的運算子，以及模型專用的 fallback 行為
效能	約 450 個 ATen 運算，以及 TorchTitan 執行 Qwen 3 8B 時專案回報的 52% MFU 結果	你的模型預熱時間、形狀、批次大小與成本
開發流程	公開範例、Notebook、除錯與基準測試工具	TPU 存取權限、Bazel 的複雜度，以及可用的可觀測性工具
目前缺口	功能覆蓋仍在持續發展	你的訓練技術堆疊是否依賴尚未針對該工作負載公開文件的功能

是否採用 TorchTPU？用這個原則判斷

如果你已經能使用 Google TPU、以 PyTorch 執行模型，並希望在訓練或推論中採用共通的框架路徑，現在就值得測試 TorchTPU。它最適合的早期使用情境，是團隊願意分析並調整張量形狀，而不是期待只改一行程式就獲得與 GPU 相當的效率。

如果你需要成熟的相容性矩陣、可預測的冷啟動行為，或工作流程依賴 TorchTPU 尚未提供的功能，那麼不妨先暫緩遷移關鍵工作負載。建議先以一個模型、一個 TPU slice 和一組固定的評估資料集分階段導入，並分別測量安裝時間、首次編譯、穩態吞吐量與輸出正確性。

TorchTPU 常見問題

TorchTPU 和 PyTorch/XLA 是同一套東西嗎？

不是。TorchTPU 是較新的原生 PyTorch 後端，設計核心是 eager execution 與 torch.compile；PyTorch/XLA 則是較早期的 TPU 整合路徑。既有的 XLA 程式碼可能需要另外評估遷移方式。

沒有 TPU 的筆電可以安裝 TorchTPU 嗎？

你可以安裝套件，但沒有相容的 Google TPU 基礎架構，就無法驗證 TPU 執行結果。官方儲存庫的實際設定前提是使用 TPU VM。

為什麼 TorchTPU 第一步執行得很慢？

第一次執行可能需要編譯計算圖並建立快取。除錯重新編譯問題時，應將後續迭代分開測量，並維持張量形狀穩定。

把 `cuda` 改成 `tpu`，就能保證有相同效能嗎？

不能。替換裝置只代表完成正確性層面的移植；比較效能前，仍需分析張量形狀、分片、運算子支援範圍與預熱時間。

TorchTPU 已經適合正式環境了嗎？

Google 表示 TorchTPU 已用於正式環境中的推論與訓練，但公開生態系元件與各項功能的成熟度不一。正式採用前，請先測試確切的模型、服務引擎與分散式模式。

在 TPU VM 上安裝指定版本的 wheel，執行 2×2 驗證程式，再用一個已完成預熱的工作負載與目前的 GPU 路徑比較。這樣才能分別看清楚程式可攜性、編譯器預熱、運算子支援範圍與硬體成本。