Cloud TPU 簡介
Tensor Processing Unit (TPU) 是 Google 專門打造的特殊應用積體電路 (ASIC)，用於加快機器學習工作負載的處理速度。如要進一步瞭解 TPU 硬體，請參閱 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
您可以透過 Compute Engine、Google Kubernetes Engine 和 Vertex AI 使用 TPU。Cloud TPU API 已淘汰。詳情請參閱「從 Cloud TPU API 遷移」。下表列出各項 Google Cloud 服務的資源。
| Google Cloud 服務 | 資源 |
|---|---|
| Compute Engine | Compute Engine 中的 TPU 資源 |
| Google Kubernetes Engine | |
| Vertex AI |
模型開發的最佳做法
如果程式是以非矩陣運算 (例如新增、重新塑形或串聯) 為主，則可能無法達到高 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
如要開始使用 Compute Engine 中的 TPU，請按照下列步驟操作：
- 如要設定專案、啟用 Compute Engine API 並取得必要權限，請按照「設定專案 Google Cloud」一文中的步驟操作。
- 如要瞭解 TPU 版本、設定、配額和價格，請參閱「規劃 Cloud TPU 資源」。
- 如要開始使用 TPU，請參閱快速入門：建立 TPU 執行個體。如要進一步瞭解如何建立 TPU，請參閱 TPU 建立總覽。
要求協助
如需協助，請與 Cloud TPU 支援團隊聯絡。 如果您有進行中的 Google Cloud 專案，請準備提供下列資訊：
- 您的 Google Cloud 專案 ID
- TPU 名稱 (如有)
- 您想提供的其他資訊