Tensor Processing Unit (TPU) 是 Google 專門打造的特殊應用積體電路 (ASIC),用於加快機器學習 (ML) 和人工智慧 (AI) 工作負載的處理速度。無論您是訓練複雜的基礎模型數週,還是執行大規模推論,TPU 都能提供可擴充的專用運算資源,並針對 JAX 和 PyTorch 等架構進行最佳化。
Cloud TPU 專為處理最嚴苛的 AI 工作負載而設計。主要福利包括:
專為矩陣運算最佳化:TPU 專為矩陣乘法單元 (MXU) 設計,可執行 ML 演算法中大量矩陣運算,效率極高。
高頻寬記憶體 (HBM):晶載高頻寬記憶體可訓練及提供較大型的模型,並有效運用較大的批次大小。
透過配量大幅擴充:TPU 晶片可以群組形式連結,稱為配量。透過切片,工作負載可擴充至數千個 TPU 晶片,以執行大規模訓練工作。
使用 TPU 的時機
TPU 最適合用於特定工作負載,例如:
- 以矩陣運算為主的模型
- 主訓練迴圈內沒有自訂 PyTorch/JAX 作業的模型
- 訓練數週或數月的模型
- 有效批量較大的大型模型
- 模型具有超大型嵌入,常見於進階排名和推薦工作負載
TPU 不適合用於下列工作負載:
- 需要頻繁分支或包含許多元素代數運算的線性代數程式
- 需要高精度計算的工作負載
- 在主要訓練迴圈中包含自訂運算的類神經網路工作負載
Cloud de Confiance by S3NS的佈建選項
您可以根據作業需求,使用下列 Cloud de Confiance by S3NS 產品 存取及佈建 TPU。
Compute Engine
透過 Compute Engine,您可以建立及管理個別 TPU VM 或切片,全面管理 TPU VM 的生命週期。Google 建議您使用 Compute Engine,而非舊版 Cloud TPU API,來佈建 TPU 資源。
詳情請參閱 Compute Engine 中的 Cloud TPU 資源。
Google Kubernetes Engine
Google Kubernetes Engine (GKE) 提供全代管的多租戶 Kubernetes 環境,可自動化調度管理大規模 AI 工作負載。GKE 支援 TPU 節點和節點集區生命週期管理,包括建立、設定及刪除 TPU VM。
詳情請參閱「GKE 中的 TPU 簡介」。
Cloud TPU
Cloud TPU API (包括 Google Cloud CLI 和 Cloud TPU 適用的 Cloud 用戶端程式庫) 已停止開發。如要佈建及管理 TPU 資源,Google 建議您根據編排和工作負載需求,使用 Compute Engine 或 GKE。
詳情請參閱「從 Cloud TPU API 遷移」。
Compute Engine 支援的 TPU 版本
Compute Engine 支援下列 TPU 版本:
- TPU7x (Ironwood)
- TPU v6e (Trillium)
- TPU v5p
如要進一步瞭解各個 TPU 版本,請參閱 TPU 機器。