Google 開源 TorchTPU，讓 TPU 成為 PyTorch 原生裝置
谷歌开源TorchTPU！TPU成PyTorch原生设备，CUDA最深的护城河开始松动
Google 開源 TorchTPU，讓 PyTorch 原生支援 TPU，並可在 Hugging Face Transformers v5 以 device_map="tpu" 部署模型。
文章梳理 TorchTPU 的執行模式、模型覆蓋與擴展測試，並指出硬件差異仍帶來調優成本，可協助理解框架兼容與實際遷移之間的落差。
編輯｜Panda
如果只能用一行代碼概括這條新聞，那就是 device_map="tpu"。
在聖何塞舉行的 PyTorch Conference 2026 上，谷歌宣佈開源 TorchTPU，這是一套讓 PyTorch 原生運行在 TPU 上的完整軟件棧。從此，在 Hugging Face Transformers v5 裏把模型部署到 TPU，只需要傳入這一個參數。
一個參數的變化，背後是谷歌在軟件路線上的一次轉向，也是英偉達 CUDA 生態近十年來遭遇的最正面的一次挑戰。
AI 芯片的競爭，從來不只是芯片本身的競爭。PyTorch 是全球使用最廣泛的 AI 框架，而它自 2016 年誕生起就與 CUDA 深度綁定：英偉達工程師長期投入，確保 PyTorch 在 N 卡上跑得最好，華爾街甚至有分析師把 CUDA 視為英偉達抵禦對手最堅固的盾牌。反觀谷歌，TPU 十年來主要服務自家的 TensorFlow 和 JAX。PyTorch 用戶想用 TPU，得經過 PyTorch/XLA 這樣一層「轉接」，付出大量額外的工程成本。
TorchTPU 改變的正是這個前提。用谷歌官方的説法，TPU 從此與 GPU、CPU 並列，成為 PyTorch 的原生選項。
項目地址：https://github.com/google-pytorch/torch_tpu
項目由谷歌與 PyTorch 的「孃家」Meta 聯合開發，vLLM、SGLang、Hugging Face、Ray、TorchTitan 等主流開源項目同步完成接入，一份 23 頁的技術報告也同時發佈在 arXiv 上。
報告標題：TorchTPU – A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU
報告地址：https://arxiv.org/pdf/2610.24161
時間點也值得關注。過去一年，Anthropic 簽下最多 100 萬顆 TPU 的使用協議，Meta 也被曝與谷歌達成數十億美元的 TPU 租用合約。硬件已經賣得出去，軟件成了最後一塊短板。現在，谷歌要親手把它補上。
不是又一個「轉接頭」
PyTorch 跑 TPU 並不新鮮。技術報告回顧，谷歌早在 2020 年就推出了基於惰性張量的 PyTorch/XLA，並罕見地直言了它的問題：需要侵入式修改源碼來劃定可編譯區域，編譯延遲嚴重，不支持 PyTorch 原生的即時執行模式，對分佈式新原語的跟進也很滯後。後來的 TorchAX 則把 PyTorch 算子逐個翻譯成 JAX 執行。在報告看來，這些方案本質上都是讓 PyTorch 去遷就 JAX 的運行方式。
TorchTPU 在 2025 年年中立項時就定下幾條原則：運行時不依賴 JAX，即時執行必須紮實，把代碼從 GPU 遷到 TPU 應當像把 .to("cuda") 改成 .to("tpu") 一樣簡單。
實現上，它基於 PyTorch 預留給第三方硬件的 PrivateUse1 接口，在 PyTorch 分發器中註冊為一個新的 C++ 後端，開發者拿到的就是普通的 PyTorch 張量，沒有子類，也沒有包裝器。
從 GPU 遷移到 TPU 前後的設備初始化代碼對比
即時執行被做成三檔：
默認的 Strict 模式逐個分發算子、與 CPU 異步執行，對應 GPU 上的常規體驗；
Debug 模式同步執行，用於排查形狀錯誤和 NaN；
第三檔 DeferAndFuse 是 TorchTPU 新引入的即時與編譯混合模式，先把一串算子攢起來，到需要取值時再自動融合成大塊計算交給 TPU，速度比 Strict 模式快 50% 以上，還能降低峯值顯存佔用。
三種模式共享一套分層編譯緩存，同樣的算子序列編譯一次即可反覆複用。
追求極致性能時，開發者可以切換到 torch.compile。這裏谷歌做了一個刻意的取捨：用 TorchDynamo 捕獲計算圖後，不走 PyTorch 默認的 Inductor 編譯器，而是把算子映射到 StableHLO，交給 XLA 生成 TPU 程序。
報告給出了理由：一是能複用為即時模式搭好的基礎設施，二是 XLA 經過多年檢驗足夠穩定，三是 Inductor 的成功很大程度上建立在 Triton 之上，而 Triton 很難移植到 TPU 架構。這句話從側面説明，PyTorch 與 GPU 的綁定，深入到了編譯器這一層。
在 TPU 上調用 torch.compile 只需指定 backend=“tpu”
PyTorch 經由 StableHLO 接入 OpenXLA 編譯棧，與 JAX、Pallas 內核共用同一條下降路徑
分佈式方面，TorchTPU 在 Process Group 層實現了 TPU 集合通信，DDP、FSDP2 和 DTensor 開箱即用，還支持 rank 0 進程多做日誌等額外工作的「不對稱」寫法，這是隻支持 SPMD 的 PyTorch/XLA 做不到的。按照 PyTorch/XLA 官方文檔的説法，TorchTPU 公開後將取代 PyTorch/XLA。
用 Stream 把集合通信卸載到 SparseCore，同時在 TensorCore 上做矩陣乘
技術報告裏的成績單
覆蓋面是第一張答卷。谷歌持續在 TPU 上測試 Hugging Face 下載量最高的約 5000 個模型，截至報告撰寫時共 4917 個、橫跨 500 多個架構族。其中 98% 能夠跑通不崩潰，在 657 個代表性架構裏，97% 的輸出與 GPU 結果數值一致。
Hugging Face 模型覆蓋情況彙總
第二張答卷是不做專門優化時的性能。谷歌只用 PyTorch 標準模塊寫了一個 30B 參數的 MuseGlimmer 模型，不調用任何模型專屬的定製內核：單顆 TPU v7 上，5120 序列長度的預填充達到每秒 2.57 萬 token，MFU 為 59%；在 4 顆 TPU v7 上訓練，MFU 為 38%。把 TorchTitan 中原本為 GPU 編寫的同款實現稍加修改搬過來，訓練 MFU 也有 36%。
第三張答卷是規模。在 Qwen3 235B MoE 模型上，64 顆 TPU v7 時 PyTorch 方案的吞吐達到谷歌自家 JAX 方案 MaxText 的約 94%。從 1024 顆芯片擴展到 8192 顆、也就是一整個 TPU v7 Pod，每個 TensorCore 的吞吐只下降約 6%，MFU 穩定在 27% 到 29% 之間，損失曲線在各個規模下幾乎重合。官方博客另稱，Qwen 3 8B 的訓練 MFU 最高可達 52%。
Qwen3 235B MoE 在 1024 到 8192 顆 TPU v7 上的擴展數據
總吞吐隨芯片數量近乎線性增長
不同規模下前 12 步的損失曲線幾乎重合
換句話説，PyTorch 在 TPU 上和「親兒子」JAX 之間，差距已經被壓到個位數百分比。
推理側的數據同樣亮眼。據報告援引的 SemiAnalysis InferenceX 結果，TPU v7 搭配 vLLM-TorchTPU 在 8K 輸入、1K 輸出的場景下，每百萬 token 成本降到 2.01 美分，比 9 月首次亮相時的 3.6 美分低了約四成。
報告還給出了幾組案例：FLUX 訓練從開箱即時模式到完整編譯優化，單步時間從 11.49 秒降至 2.78 秒；Meta 的 DLRM v2 推薦模型在 32 顆 TPU v7x 上與谷歌高度優化的 JAX 基線完全持平。按報告的説法，TorchTPU 已經在部分谷歌雲大客戶的生產環境中承擔推理和訓練任務。
一次「全家桶」式的開源
比技術本身更有分量的，是生態的完整度。
推理側，與 Inferact 合作的 vLLM 後端、與 RadixArk 合作的 SGLang 後端都保留了上游原有的調度器和 API 服務器，現有推理棧無需改動；谷歌還聯合 Red Hat、IBM 發佈 llm-d 1.0，在 GKE 上讓 GPU 與 TPU 共用同一個推理入口。
訓練側，TorchTitan 提供稠密模型、MoE 和擴散模型的預訓練配方，強化學習框架 VeRL 合入了 TPU 後端，Ray 支持把同一工作流的不同階段分別調度到 CPU、GPU 或 TPU。
端側，前身為 TensorFlow Lite 的 LiteRT 正式加入 PyTorch 生態。TorchTPU、vLLM-TorchTPU、SGLang-TorchTPU 和 TorchTitan 四個核心倉庫同步開源。
站台方的表態同樣值得玩味。Meta PyTorch 工程總監 Jana van Greunen 直言，Meta 不能被單一硬件架構束縛，TorchTPU 已經支撐起其 LLM、廣告和推薦系統的關鍵生產場景。要知道，Meta 可是英偉達的超級大客戶。
從內部代號到 GitHub
回看 TorchTPU 的路徑，節奏相當緊湊。2025 年年中，項目在谷歌內部啓動 ；同年 12 月，路透社首次披露這一代號，稱谷歌正與 Meta 密切合作，並在考慮開源部分代碼。2026 年 4 月，谷歌開發者博客公開技術架構，項目負責人在巴黎 PyTorch Conference Europe 上登台。9 月，TorchTPU 進入私有 Beta，並在 SemiAnalysis 的推理基準中首次亮相。到今天，整套棧正式開源。
從「考慮開源部分代碼」到四個核心倉庫全部公開、各大項目上游合入，谷歌給出的答案比外界預期的更徹底。
CUDA 護城河會被填平嗎
一個更進一步的問題是：框架層的壁壘被拆除後，英偉達還剩下什麼？
首先要看到，TorchTPU 抹平的是框架層，而不是硬件差異本身。報告自己也承認，只改一行 .to("tpu") 拿不到峯值性能，因為模型的形狀、分片和算子選擇可能本就偏向 GPU；TPU 的 TensorCore 本質上是單線程處理器，形狀要足夠寬才能餵飽矩陣單元，這與 GPU 靠並行跑大量小內核的思路截然不同。
一些細節差異同樣需要遷移者留意：TorchTPU 默認採用 BF16 乘法加 FP32 累加，而 PyTorch 在 CUDA 上默認是完整 FP32；默認的集合通信編譯方式會帶來 30% 到 40% 的端到端損失，需要額外標註來優化；每顆 TPU 只能被一個進程訪問，強化學習中訓練與推理共置的做法暫時無法開箱支持。「能跑」和「跑得好」之間，仍然隔着調優的工作量。
PyTorch 在 CUDA 與 TPU 上的數值精度對比
其次，這並非谷歌一家的動作。PyTorch Conference 2026 上，AWS 也在講 Trainium 走向原生 PyTorch 的歷程，其 Neuron SDK 已於今年停止支持用 PyTorch/XLA 做訓練，轉向原生的 TorchNeuron。
TorchNeuron 的高層架構圖
PyTorch 正在變成一層事實上的硬件抽象層，PrivateUse1 則成了各家芯片接入的標準插槽。當所有硬件都能以原生方式跑同一份 PyTorch 代碼，競爭焦點就會從「生態能不能用」轉向「每美元能換來多少吞吐」，這恰恰是谷歌在發佈稿裏反覆強調的指標。
當然，變量依然存在。報告中的大部分性能數據由谷歌自己測得，還需要社區在更多模型和場景中獨立驗證；TPU 能否以足夠的產能和價格觸達普通開發者，同樣決定着這次開源的最終影響。谷歌也表示，未來 6 到 12 個月的重點是生產環境加固、動態形狀支持以及面向下一代 TPU 的 FP8、FP4 量化。
但方向已經清晰。過去十年，選擇 PyTorch 幾乎就等於選擇了 CUDA。從今天起，這道等式不再成立。
參考鏈接
https://cloud.google.com/blog/topics/ai-infrastructure/torchtpu-is-now-open-source-run-native-pytorch-on-tpus
https://developers.googleblog.com/torchtpu-running-pytorch-natively-on-tpus-at-google-scale/
https://www.denvergazette.com/2025/12/17/google-works-to-erode-nvidias-software-advantage-with-metas-help/
https://pytorchconferenceeu2026.sched.com/event/2Ikir
來源：新浪科技 · finance.sina.com.cn