Cloud TPU 簡介

Tensor Processing Unit (TPU) 是 Google 專門打造的特殊應用積體電路 (ASIC)，用於加快機器學習工作負載的處理速度。如要進一步瞭解 TPU 硬體，請參閱「TPU 架構」。Cloud TPU 是一項網路服務，可讓 TPU 在 Google Cloud上提供可擴充的運算資源。

TPU 會使用專為執行機器學習演算法中常見的大型矩陣運算而設計的硬體，更有效率地訓練模型。TPU 具有晶載高頻寬記憶體 (HBM)，可讓您使用較大的模型和批次大小。TPU 可連結成名為「切片」的群組，藉此在不需變更程式碼的情況下，擴充工作負載。

在 TPU 上執行的程式碼必須由加速器線性代數 (XLA) 編譯器編譯。XLA 是一種及時編譯器，可採用 ML 架構應用程式所發出的圖表，並將圖表的線性代數、損失和梯度元件編譯為 TPU 機器碼。程式的其餘部分則在 TPU 主機上執行。XLA 編譯器是 TPU 主機機器上執行的 TPU VM 映像檔的一部分。

如要進一步瞭解 Tensor Processing Unit，請參閱「如何看待 TPU」。

何時使用 TPU

Cloud TPU 最適合用於特定工作負載。在某些情況下，您可能需要在 Compute Engine 執行個體中使用 GPU 或 CPU 執行機器學習工作負載。一般來說，您可以依據下列準則，決定哪些硬體最適合您的工作負載。

CPU

  • 迅速建立需要最高彈性的原型
  • 訓練時間不長的簡單模型
  • 具備小型有效批量的模型
  • 包含許多 以 C++ 編寫的自訂 TensorFlow 運算的模型
  • 受限於主機系統可用 I/O 或網路頻寬的模型

GPU

  • 具備大量自訂 PyTorch/JAX 運算，且至少必須於 CPU 上執行部分的模型
  • TensorFlow 運算無法在 Cloud TPU 使用的模型 (請參閱可用的 TensorFlow 運算清單)
  • 有效批量較大的中型至大型模型

TPU

  • 以矩陣運算為主的模型
  • 在主要訓練迴圈中沒有自訂 PyTorch/JAX 運算的模型
  • 訓練數週或數月的模型
  • 有效批量較大的大型模型
  • 模型中含有超大型嵌入資料，這類模型常見於進階排名和推薦工作負載

Cloud TPU 不適合用於下列工作負載：

  • 需要頻繁分支或包含許多元素代數運算的線性代數程式
  • 需要高精度運算的工作負載
  • 在主訓練迴圈中包含自訂運算的類神經網路工作負載

Google Cloud中的 TPU

您可以透過 Cloud TPU VM、Google Kubernetes Engine 和 Vertex AI 使用 TPU。下表列出各項 Google Cloud服務的資源。

Google Cloud service 資源
Cloud TPU 開始使用 Cloud TPU VM
Google Kubernetes Engine

關於 GKE 中的 TPU

使用 TPU 在 GKE 中執行 Ray
Vertex AI

使用 TPU 在 Vertex AI 上訓練

在 Vertex AI 上使用 TPU 進行線上預測

模型開發的最佳做法

如果程式是以非矩陣運算 (例如新增、重新塑形或串聯) 為主，則可能無法達到高 MXU 使用率。以下準則可協助您選擇及建構適合使用 Cloud TPU 的模型。

版面配置

XLA 編譯器會執行程式碼轉換，其中包括將矩陣乘法並排顯示為較小區塊，以便有效率地在矩陣單元 (MXU) 執行運算。XLA 編譯器會使用 MXU 硬體結構 (128x128 的收縮自激陣列) 和 TPU 記憶體子系統設計 (偏好 8 的倍數維度)，以提高平鋪效率。

因此特定版面配置更容易進行並排顯示，其他版面配置則需要先「重新塑形」才能並排顯示。在 Cloud TPU 上，重新調整形狀的運算通常會受到記憶體限制。

形狀

XLA 編譯器會針對第一個批次及時編譯機器學習圖。如果任何後續批次擁有不同形狀，模型就無法運作。(如果每次都將圖重新編譯，則形狀改變速度太慢)。因此，任何具有動態形狀張量的模型都不適合使用 TPU。

填充

高效能的 Cloud TPU 程式可將密集運算分割成 128x128 的區塊。如果矩陣運算無法佔用整個 MXU，編譯器會以零填補張量。填充有兩項缺點：

  • 以零填充的張量會使 TPU 核心使用率過低。
  • 填充會增加張量所需的晶載記憶體儲存量，在極端情況下可能會造成記憶體不足的錯誤。

雖然填充是由 XLA 編譯器於必要時自動執行，但您可利用 op_profile 工具決定執行的填充量。您可以選擇適合 TPU 的張量維度，避免填充。

維度

選擇適當的張量維度，能大幅協助 TPU 硬體發揮最高效能，MXU 特別是如此。XLA 編譯器會嘗試使用批量或特徵維度，盡量充分地運用 MXU，因此這兩者之中至少要有一個為 128 的倍數，否則編譯器會將其中一項填充至 128。理想情況下，批次大小和功能維度應為 8 的倍數，這樣才能從記憶體子系統中擷取高效能。

開始使用 Cloud TPU

  1. 設定 Google Cloud 帳戶
  2. 啟用 Cloud TPU API
  3. 授權 Cloud TPU 存取 Cloud Storage 值區
  4. 在 TPU 上執行基本計算
  5. 在 TPU 上訓練參考模型
  6. 分析模型

尋求協助

如需協助，請與 Cloud TPU 支援團隊聯絡。如果您有 Google Cloud 有效專案，請準備提供下列資訊：

  • 您的 Google Cloud 專案 ID
  • 您的 TPU 名稱 (如果有的話)
  • 您想提供的其他資訊

後續步驟

是否要進一步瞭解 Cloud TPU？以下資源可能有所幫助：