2026 年 10 月 20 日，Google Cloud 在 PyTorch Conference 2026 宣布 TorchTPU 正式開源，程式碼存放於 GitHub 的 google-pytorch/torch_tpu 專案，採用 Apache 2.0 授權。TorchTPU 由 Google 與 Meta 的 PyTorch 團隊共同開發，把 TPU 變成 PyTorch 的原生執行目標，與 GPU、CPU 並列[1]。官方指出，現有 PyTorch 模型只需把 model.to("cuda") 改成 model.to("tpu")，即可在 TPU 上運行。
捨棄延遲張量路線
Google 早於 2020 年推出 PyTorch/XLA，以延遲張量（lazy tensor）方式把 PyTorch 編譯到 TPU[2]。官方承認這個方案要求開發者在原始碼中標明可編譯區域，編譯延遲嚴重，亦不支援 PyTorch 的命令式（eager）執行模式，難以追上社群演進的速度。TorchTPU 改為建基於 PyTorch 的 Dispatcher 與 Torch Dynamo，支援約 450 個 Aten 算子，覆蓋主要深度學習架構所需的前向與反向運算，並提供名為 DeferAndFuse 的 eager 模式[2]。開發者可以先用 eager 模式除錯，再以 torch.compile 取得正式環境的效能；自訂算子則可透過 JAX/Pallas 或 Helion DSL 編寫，執行期不依賴 JAX。
生態整合與分散式支援
分散式原語實作於 Process Group API 層，繼承 c10d::Backend，因此 DTensor、FSDP、SimpleFSDP 等框架毋須改動即可使用，集體通訊以 StableHLO 算子表達[2]。推論方面，vLLM 與 SGLang 均有對應的 TorchTPU 分支；訓練由 TorchTitan 支援，並可配合 Ray 編排。Hugging Face 同步更新，Transformers v5 只要傳入 device_map="tpu" 就能把模型部署到 TPU，Diffusers 亦原生支援 eager、torch.compile 與張量並行[1]。Google 同時發布 llm-d 1.0，與 Red Hat、IBM 等社群合作，在 Kubernetes 上以同一端點調度 TPU 與 GPU 資源[1]。
實測數據與實際限制
官方公布的測試中，Llama 3 405B 在 1,024 顆 TPU v7x 上以 BF16 及 8K 上下文訓練，達到每顆晶片每秒 500 個 token、MFU 超過五成；Qwen3 235B 混合專家模型在 8,192 顆 TPU v7x 上、4K 序列長度下達到 28% MFU[2]。在 64 顆 TPU v7 的比較中，TorchTitan 錄得 1,828 TPS，JAX 的 MaxText 則為 1,937 TPS[2]，兩者差距已相當有限。不過官方明確指出，改一行程式碼只代表模型「跑得起身」，要達到理想效能仍須按 TPU 的形狀、切分方式與算子選型調整，開發者不宜把單行遷移當成終點[1]。
部署前要注意的事
TorchTPU 目前支援 Python 3.11 至 3.15，並已在 Google Cloud 部分大型客戶的推論與訓練工作負載中投入使用[1][2]。對打算在 TPU 上跑現有 PyTorch 專案的團隊而言，較穩妥的做法是先以 eager 模式驗證數值結果與收斂曲線，再逐段啟用 torch.compile，並以 TorchTitan 或 vLLM 的對應分支作為參考實作。若模型依賴 FlexAttention 等尚未支援的算子，官方文件表明需要先以標準的 scaled_dot_product_attention 替代，遷移成本不能低估。
參考資料
[1] TorchTPU is now open source: Run native PyTorch on TPUs（2026年10月20日）