Tensor Processing Unit (TPU) 是 Google 開發的客製化特殊應用積體電路 (ASIC),專門用來加速機器學習 (ML) 和人工智慧 (AI) 工作負載。無論是訓練複雜的基礎模型數週,還是執行大規模推論,TPU 都能提供可擴充的專用運算資源,並針對 JAX 和 PyTorch 等框架進行最佳化。
Cloud TPU 專為處理最嚴苛的 AI 工作負載而設計。主要優點包括:
針對矩陣運算最佳化:TPU 專門設計了矩陣乘法單元 (MXU),能以極高效率執行機器學習演算法所需的大量矩陣運算。
高頻寬記憶體 (HBM):晶片內建高頻寬記憶體可讓您訓練及提供較大的模型,並有效運用較大的批次大小。
透過配量大規模擴充:TPU 晶片可以分組連結,稱為配量。配量可讓工作負載擴充至數千個 TPU 晶片,用於大規模訓練工作。
使用 TPU 的時機
TPU 經過最佳化調整,適合特定工作負載,例如:
- 以矩陣運算為主的模型
- 主要訓練迴圈中沒有自訂 PyTorch/JAX 作業的模型
- 訓練數週或數月的模型
- 有效批量較大的大型模型
- 模型具有超大型嵌入,常見於進階排名和推薦工作負載
TPU 不適用於下列工作負載:
- 需要頻繁分支或包含許多元素代數運算的線性代數程式
- 需要高精度算術的工作負載
- 在主要訓練迴圈中包含自訂作業的類神經網路工作負載
Google Cloud的佈建選項
您可以根據作業需求,使用下列 Google Cloud 產品存取及佈建 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 簡介」。
批次
Batch 是一項全代管服務,可讓您排定、排入佇列及執行容器化或指令碼型工作負載,這些工作負載會在 Compute Engine TPU VM 上執行。
詳情請參閱「建立及執行使用 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 機器」。