Cloud TPU 簡介

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

TPU 採用專為執行大型矩陣運算設計的硬體，可更有效率地訓練模型，這類運算通常出現在機器學習演算法中。TPU 具有晶載高頻寬記憶體 (HBM)，可供您使用較大的模型和批次大小。TPU 可連結至名為「切片」的群組，以便擴充工作負載，且幾乎不需要變更程式碼。

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

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

何時使用 TPU

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

CPU

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

GPU

  • 具備大量自訂 PyTorch/JAX 運算，且至少必須於 CPU 上執行部分的 PyTorch/JAX 運算的模型
  • 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 服務 資源
Cloud TPU 開始使用 Cloud TPU VM
Google Kubernetes Engine

關於 GKE 中的 TPU
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？以下資源可能有所幫助：