Google 與 Meta 開源 TorchTPU，為 Google TPU 提供原生 PyTorch 支援

TorchTPU 是 Google 與 Meta 合作開發的開源項目，為 Google TPU 提供原生 PyTorch 後端。它只需極少的程式碼修改，便可讓 PyTorch 工作負載在 TPU 硬件上執行，並支援即時及編譯執行、分散式運算，以及硬件專屬最佳化。

Google 於 2020 年開始透過 PyTorch/XLA 讓 PyTorch 工作負載在 TPU 上執行；PyTorch/XLA 利用了 JAX、XLA 及惰性張量。Google 指出，這種做法有明顯不足：

最重要的是，PyTorch/XLA 需要侵入式修改原始碼，以明確界定可編譯區域；它的編譯時間延遲嚴重，不支援 PyTorch 原生的命令式（「eager」）執行模式，而且在支援 PyTorch 生態系統中新興的分散式運算原語方面落後。

TorchTPU 採用截然不同的做法，將 eager 執行列為首要目標，讓現有 PyTorch 工作負載只需少量程式碼修改便可執行。

具體而言，從 GPU 遷移至 TPU，需要將 cuda 裝置替換為 tpu，並將 nccl 分散式後端替換為 tpu_dist。開發者亦可使用新的 torch.accelerator.current_accelerator API 識別可用的加速器，如以下程式碼片段所示。

import torch
import torch.distributed as dist

def main():
    device = torch.device("cuda") # -> device = torch.accelerator.current_accelerator
    dist.init_process_group(backend="nccl") # -> dist.init_process_group()
    rank = dist.get_rank()
    word_size = dist.get_world_size()
    ...
    model = model.to(device)
    ...

TorchTPU 支援三種 eager 執行模式：strict、debug 和 DeferAndFuse。後者是全新的 eager／編譯混合執行模式。Google 表示，這種模式「速度可比 Strict Eager 快 50% 或以上，同時降低 HBM 峯值記憶體用量」。

為測試 TorchTPU 能否在不作任何修改的情況下執行現有 PyTorch 模型，Google 建立了一條整合管線，在 v5e 起的各代 TPU 硬件上，以 eager 和編譯模式評估約 5,000 個 Hugging Face 模型，涵蓋 500+ 種架構。

我們採用穩健的測試方法，以語料分層和均衡的架構選取為基礎。我們的做法透過雙層選取流程，兼顧實際採用情況和廣泛的架構多樣性，避免重複測試細微變體，同時確保真正用於生產環境的權重得到全面評估。

為驗證 TorchTPU 的效能，Google 使用標準 PyTorch 實作了一個 300 億參數的 Muse Glimmer 模型。在單個 TPUv7 晶片上，序列長度為 5120 時，推理吞吐量達到 25.7k token/s；序列長度為 10,240 時，則達到 22.4k token/s。訓練方面，使用四個 TPUv7 晶片時，序列長度分別為 5120 和 10,240，吞吐量達到 5.6k token/s 和 4.6k token/s。

Google 亦使用 Qwen3 MoE 235B 進行訓練工作負載實驗，將規模由 1,024 個 TPUv7 晶片擴展至 8,192 個。推理方面，Google 表示，使用 vLLM 處理大型 MoE 工作負載時，每美元效能具競爭力。

對於這項公佈，Hacker News 用戶 tingkuanpei 指出：

非常有意思的工作。Google 早就應該做這件事。PyTorch 的生態系統實在非常豐富。

其他用戶指出，Google 除了透過 GCP 提供 TPU，亦會將 TPU 出售給「用於本地部署的企業客戶」，其中包括 Anthropic；Meta 亦可能在未來跟進。

TorchTPU 仍處於早期階段，因此持續演進。該專案的發展路線圖涵蓋動態形狀、編譯效能、運算子覆蓋範圍、強化學習及量化等工作。如欲瞭解詳情，請務必閲讀原始公告，當中更全面地介紹了 TorchTPU 的架構和範圍。

關於作者

Sergio De Simone

顯示更多顯示較少