在 Cloud TPU 上執行 PyTorch

PyTorch on Cloud TPU 可讓您使用標準 PyTorch 語意、分散式 API 和現有的生態系統工具 (包括 Hugging Face、vLLM、SGLang、llm-d、TorchTitan、TorchRec 和 Ray)，在 Cloud TPU 硬體上訓練、微調及提供開放原始碼模型。

這個堆疊的基礎是 TorchTPU，這是與 Meta 的 PyTorch 團隊共同開發的開放原始碼 PyTorch 原生後端。TorchTPU 提供程式庫、架構和工具，涵蓋機器學習生命週期的每個階段，包括核心執行階段執行、自訂核心、正式環境推論、大規模訓練和叢集自動化調度管理。

TorchTPU 連結至 Hugging Face、TorchTitan、vLLM、SGLang、Ray 和 llm-d，並搭配 torch.profiler、DTensor、TorchRec 和 Pallas。

開始使用

TorchTPU：Cloud TPU 的 PyTorch 原生後端

TorchTPU 提供裝置後端、圖形編譯器整合，以及 Cloud TPU 上 PyTorch 的自訂核心 API。TorchTPU 是與 PyTorch 團隊合作開發，可讓您在 eager 模式下以互動方式偵錯模型、使用 torch.compile 編譯模型程式碼，以及使用 torch.distributed (tpu_dist) 跨切片擴充。

如要安裝 TorchTPU 並在 Cloud TPU 上執行張量運算，請按照下列步驟操作：

pip install torch_tpu

import torch

device = torch.device("tpu")

x = torch.randn((4, 4), device=device)
y = torch.ones((4, 4), device=device)
result = torch.matmul(x, y).cpu()

print(f"Executed on device: {x.device}")
print(f"Result shape on CPU: {result.shape}, sum: {result.sum():.4f}")

動機

雖然 Cloud TPU 過去一直與 JAX 和 TensorFlow 搭配使用，但先前在 TPU 上執行 PyTorch 時，需要延遲張量圖記錄 (PyTorch/XLA)，這通常代表需要新增程式碼來標記可編譯區域。TorchTPU 預設會使用 PyTorch 的標準單一程序/裝置執行模型，積極執行作業，因此現有的工作負載和程式庫不需重寫結構即可執行。

設計

TorchTPU 為 PyTorch 帶來 PyTorch 原生 TPU 支援，並著重於三個設計目標：

  • 標準 PyTorch 語意：每個裝置執行一項程序，並保留 PyTorch 張量檢視和就地變異，不依賴 JAX 執行階段。
  • 預設為 Eager Execution：標準 PyTorch 模型會立即在 Eager 模式下執行，方便您進行互動式偵錯和逐步檢查。
  • 編譯器和核心最佳化：除了逐一運算子的 Eager Execution 之外，TorchTPU 還提供 Eager 模式下的運算子融合、torch.compile 整合動態形狀邊界以減少重新編譯，以及自訂核心 API。

在圖形編譯方面，torch.compile 會使用 TorchDynamo 擷取 PyTorch 運算圖，並透過 OpenXLA 和 StableHLO 將其編譯為 TPU 機器碼。集體運算會在 torch.distributed 程序群組層級 (tpu_dist) 執行。

主要優點

  • 運算子、模型和 Hugging Face 支援：TorchTPU 支援正向和反向傳遞中的標準 PyTorch (ATen) 運算子，並持續針對熱門 Hugging Face 模型進行測試。您可以傳遞 device_map="tpu"，在 Transformers v5 中部署模型；使用 eager 模式、torch.compile 和張量平行處理執行 Diffusers 管道；以及在 TensorBoard 和 XProf 中使用 torch.profiler 分析工作負載。
  • 從單晶片到多配量的擴充：tpu_dist 後端支援 DistributedDataParallel (DDP)、FullyShardedDataParallel (FSDP2) 和 DTensor。此外，torch.tpu.distributed.topology_aware_mesh 會將平行處理維度對應至實體 TPU 互連。
  • 可設定的 Eager Execution 模式：提供三種 Eager Execution 模式：
    • Strict (預設)：使用標準命令式 PyTorch 語意，個別且非同步調度運算子。
    • Debug：同步執行，並檢查張量邊界。
    • DeferAndFuse：融合 eager 運算子序列，減少高頻寬記憶體 (HBM) 流量。
  • 完整圖表編譯和快取：torch.compile(backend="tpu") 會將圖表編譯成 TPU 執行檔，並在程序、主機和工作重新啟動時快取這些檔案。
  • 自訂 TPU 核心：您可以呼叫自訂 Pallas 核心 (包括直接在晶片之間通訊的核心)，做為標準 PyTorch 函式、使用 Tokamax 程式庫中的預先建構核心，或使用 Helion 編寫可攜式圖塊核心。

如需更多資訊，請參閱 TorchTPU 說明文件、在 PyTorch on TPU 討論論壇中提問，或在 GitHub 回報問題。

正式環境推論與提供

如要部署至正式環境，Cloud TPU 上的 PyTorch 可與開放原始碼推論引擎和 Kubernetes 原生分散式服務基礎架構整合。

vLLM-TorchTPU 和 SGLang-TorchTPU：推論引擎

vLLM-TorchTPU 和 SGLang-TorchTPU 是 vLLM 和 SGLang 推論引擎的 TorchTPU 後端。這兩個後端都執行開放式模型，例如 Qwen、DeepSeek、GLM、Kimi、GPT-OSS 和 Gemma 4。詳情請參閱網誌文章「Announcing New vLLM and SGLang Inference Engines on TPU」(宣布推出 TPU 適用的全新 vLLM 和 SGLang 推論引擎)。

動機

如要提供推論、多模態和工具使用工作負載，通常需要自訂模型定義、推測解碼方法和量化配置。在 TorchTPU 建構 vLLM 和 SGLang 後端，即可直接在 PyTorch 自訂及偵錯服務堆疊。

設計

這兩個後端都保留標準 vLLM 和 SGLang 排程器、KV 快取管理工具和 API 伺服器，同時在底層使用 TorchTPU 執行和自訂 Pallas 運算子。大多數模型會執行標準 PyTorch 模型定義，而運算密集區塊則會在編譯圖中叫用 Pallas 核心，因此可獨立更新硬體專屬核心。

主要優點

  • 平行處理、量化和推測解碼：支援跨主機的張量、資料、專家、管道和情境平行處理，以及多模態 (圖片和影片) 和嵌入輸入。支援的推測解碼方法包括多詞元預測 (MTP)、EAGLE-3、DFlash、DSpark 和 N-gram 比對。支援的量化格式包括 8 位元 (FP8) KV 快取，以及 8 位元 (FP8) 和 4 位元 (MXFP4 和 NVFP4) 浮點模型檢查點。
  • 雙層前置字元快取、工具呼叫和結構化輸出：自動前置字元快取 (包括混合注意力機制和遞迴模型) 會將有效前置字元保留在 HBM 中，並將舊前置字元卸載至主機記憶體，避免在多輪要求中重新計算。內建工具呼叫和結構化輸出功能。
  • 使用 TPU Sync 重疊 KV 快取傳輸：支援分散式預先填充和解碼服務 (將提示詞處理和詞元生成作業分開至不同的 worker 集區)，可跨主機和主機 KV 卸載，使用 TPU Sync。TPU Sync 是開放原始碼 C++ 程式庫，可直接在背景執行緒讀取 PyTorch 裝置緩衝區，不需透過主機 CPU 記憶體複製資料。

如要進一步瞭解 vLLM-TorchTPU，請參閱 vLLM-TorchTPU 存放區、在 vLLM Slack 頻道提問，或在 GitHub 回報 vLLM-TorchTPU 問題。

如要進一步瞭解 SGLang-TorchTPU，請參閱 SGLang-TorchTPU 存放區，或在 GitHub 回報 SGLang-TorchTPU 問題。

llm-d：Kubernetes 原生分散式推論

llm-d 是開放原始碼的雲端原生運算基金會 (CNCF) 沙箱框架，用於在 Kubernetes 上執行分散式 LLM 推論，由 Red Hat、IBM、 Google Cloud，以及超過 500 名貢獻者和 35 個組織組成的社群共同開發。

動機

雖然單一執行個體推論引擎會管理每個副本的執行作業，但正式環境機群必須處理流量高峰、多租戶優先順序、機群範圍內容共用，以及零停機時間推出作業。llm-d 提供叢集層級轉送和作業層，可協調推論伺服器機群。

設計

llm-d 會在 Kubernetes 上的模型伺服器 (例如 vLLM 和 SGLang) 前端執行，並提供下列功能：

  • 根據前置字元快取狀態和副本負載轉送要求。
  • 協調預先填入和解碼工作站集區。
  • 支援 MoE 分區和非同步批次推論。
  • 管理裝置記憶體、主機記憶體和儲存空間的分層 KV 快取卸載。

主要優點

  • 多輪工作階段路由：將多輪對話路由至保留快取脈絡的副本。輪流優先排程和有狀態的 Responses API 支援，可讓代理重複使用各輪對話的背景資訊，減少重新運算。
  • 多租戶服務品質：優先順序和公平流量控制功能，可讓對延遲敏感的要求和突發代理工作負載共用機群，不會發生資源匱乏的情況。
  • 機群作業和滾動式更新：提供 v1 設定 API、路由器高可用性，以及協調式滾動式更新，將預先填入和解碼副本視為單一邏輯單元管理。
  • 在 GKE 上提供原生 TorchTPU：您可以在一個 llm-d 端點後方，與其他加速器一起提供 Cloud TPU，並使用標準 Kubernetes 工具管理 vLLM-TorchTPU 和 SGLang-TorchTPU 副本。

如需更多資訊， 請在 llm-d Slack 頻道中提問，或在 GitHub 回報問題。

訓練、微調和領域程式庫

TorchTPU 提供參考訓練方法和特定領域程式庫，並針對 TPU 硬體進行調整，可預先訓練、微調及調整大型模型。

在 TPU 上使用 TorchTitan：大規模訓練參考資料

TPU 上的 TorchTitan 擴充了 TorchTitan，這是 PyTorch 模組化參考平台，用於生成式 AI 訓練，具備 TPU 分區傳遞、自訂核心和訓練配方。詳情請參閱網誌文章「Integrating TorchTPU into TorchTitan and Optimizing Models for TPU」(將 TorchTPU 整合至 TorchTitan 並針對 TPU 最佳化模型)。

動機

如要跨多主機 TPU 切片擴充訓練，必須在晶片互連網路中分割參數，並重疊通訊與運算。Google 和 Meta 共同開發了 TorchTitan 內部的 TPU 分區和編譯階段，因此您不必設計自訂 TPU 分區策略，也不必手動排定通訊。

設計

模型是以單一裝置 PyTorch 編寫，不含加速器專屬程式碼。TorchTitan 會根據設定檔設定模型層，因此您可以使用自訂 Pallas 核心 (例如 Splash Attention) 取代個別層，而不必編輯模型定義或訓練迴圈。

TorchTitan 的 SimpleFSDP 實作會將參數分片作業直接放入編譯圖中，因此 XLA 編譯器可以將裝置通訊與運算重疊，而全步驟圖編譯則會將整個訓練步驟擷取為單一 XLA 程式。如果是多切片執行作業，混合式 sharding (HSDP) 會透過晶片間互連 (ICI) 在每個切片內 sharding，並透過資料中心網路 (DCN) 跨切片複製。這種配置會將跨切片流量限制為每個步驟減少一個梯度。

主要優點

  • 預先訓練和強化學習後訓練配方：包含稠密 LLM (Llama 3、Qwen 3 和 MuseGlimmer 30B)、混合專家 (MoE) 模型 (Qwen 3 235B，已在多達 8,192 個 TPU7x 晶片上驗證) 和擴散轉換器 (FLUX.1) 的參考配方。如果是強化學習訓練後調整 (例如群組相對策略最佳化，或 GRPO)，VeRL (預先發布版) 會將參數更新委派給 TorchTitan，並將更新後的權重推送至服務引擎，以產生訓練軌跡 (展開)。
  • 預先設定的編譯器和平行處理設定：配方設定包含預先設定的編譯器和平行處理設定，因此您可以透過設定檔調整模型 FLOP 使用率 (MFU)，而不必變更程式碼。
  • 新架構入門指南：Hugging Face Transformers 建模後端會透過 TorchTitan 訓練器，執行 Hugging Face Hub 模型檢查點 (包括不含專家平行處理的 MoE 模型)。針對完整模型移植，TorchTitan 提供逐步入門技能，協助工程師和程式設計代理。
  • 獨立元件重複使用：您可以端對端採用 TorchTitan，或將個別模組、自訂運算子和編譯器旗標匯入現有的訓練程式碼集。

如要瞭解詳情，請參閱 TorchTitan on TPU 入門指南、在 PyTorch 討論論壇中提問，或在 GitHub 回報問題。

TorchRec：推薦系統

TorchRec 是 PyTorch 領域程式庫，提供分散式嵌入表格和建議模型的分區原始物件。在 Cloud TPU 上，TorchRec 會將稀疏嵌入表格查詢作業卸載至專屬的 TPU SparseCore 處理器，同時在 TPU TensorCore 上執行密集矩陣運算。

如要進一步瞭解 TorchRec，請參閱 TorchRec 說明文件、在 PyTorch on TPU 討論論壇中提問，或在 GitHub 回報問題。

叢集自動化調度管理和開發人員工具

為擴充多節點管道並自動化模型上線流程，TorchTPU 生態系統與 Ray 整合，並提供可存取 TPU 說明文件的程式碼代理擴充功能。

Ray：分散式訓練、提供和強化學習

Ray 使用 Ray Train、Ray Serve、Ray Data、Ray Core 和 verl 等強化學習框架，在叢集間擴充 Python 和 AI 工作負載。

動機

多配量訓練和多級式管道 (例如結合資料預先處理、訓練和生成推出作業的 RL 工作流程) 需要自動化配量佈建、拓撲探索和容錯 worker 協調。torchrun 等獨立啟動公用程式不會佈建 TPU 配量，也不會協調多階段 actor 管道，因此 Ray 會管理 TorchTPU 工作負載的叢集範圍排程和配置。

設計

提交工作時，只要指定 TPU 世代、拓撲和 worker 數量，Ray 就會保留完整的 TPU 切片，以便同時排定切片中的所有 worker、找出晶片互連位址，並為每個 TPU 裝置啟動一個 PyTorch worker。Ray Train 的 PyTorch 後端會在啟動工作前，自動初始化 tpu 裝置和 tpu_dist 程序群組。

主要優點

  • 拓撲相符的配置：根據實體 TPU 世代和配量版面配置，在主機之間配置 worker。
  • 自動復原故障：在 GKE 上跨多個 TPU 切片執行分散式訓練工作，並在 worker 或切片故障時自動復原。
  • 直接預先擷取資料至 TPU：Ray Data 的 iter_torch_batches 方法會將批次資料預先擷取為 PyTorch 張量，直接載入 TPU 裝置記憶體。
  • 多主機推論自動化調度管理：管理 vLLM-TorchTPU 和 SGLang-TorchTPU 的多主機服務部署作業。

TPU agent suite：(預先發布版) TPU 程式設計代理

TPU agent suite 是程式開發輔助工具的開放原始碼擴充功能，包括 Claude Code 和 Google Antigravity，未來還會支援更多程式開發代理。這項工具專門用於 TPU v6e 和 TPU7x 上的 Transformer 推論 (密集 LLM 和 MoE 模型)。

動機

一般用途的程式開發輔助通常缺乏 TPU 硬體拓撲、XLA 編譯器旗標、XProf 遙測和 Pallas 核心限制的最新背景資訊。TPU agent suite 可讓 AI 助理存取官方 TPU 說明文件和結構化工作流程。這有助於避免無效的標記建議，並自動執行疊代剖析和基準測試。

設計

這套工具會在現有的程式開發輔助中，加入 TPU 知識庫和六個特定任務子代理。路由器代理會將每個提示委派給適當的子代理，不會自行執行指令。在 TPU VM 或 GKE 叢集上執行任何指令，都必須經過您的核准，而且調整迴圈會執行有限次數的記錄嘗試，然後才會傳回結果。

主要優點

  • 模型導入和準確度驗證：TPU 模型導入代理會解決不支援的運算子、選取服務路徑 (vLLM-TorchTPU 或標準 PyTorch)，並將逐層輸出內容與 CPU 參考模型進行比較，驗證數值準確度。
  • 當機和瓶頸診斷：TPU 診斷代理會分析當機記錄、檢查 XProf 設定檔和 XLA 編譯器追蹤記錄，並找出執行階段瓶頸。
  • 核心、編譯器和分區最佳化：TPU 核心開發代理可協助自訂 Pallas 核心，TPU 效能最佳化代理可基準測試 XLA 編譯器旗標組合，TPU 分區代理則可規劃多切片平行處理版面配置。
  • 查詢文件並提供引用來源：TPU 知識庫代理可回答技術問題，並明確引用官方 TPU 文件和教學課程，也會說明未涵蓋的主題。

如要瞭解詳情，請參閱 TPU agent suite 說明文件、在 GitHub 討論區分享意見和提問，或在 GitHub 回報問題。

後續步驟