瞭解 TPU

Tensor Processing Unit (TPU) 是 Google 開發的客製化特殊應用積體電路 (ASIC)，專門用來加速機器學習 (ML) 和人工智慧 (AI) 工作負載。無論您是訓練複雜的基礎模型，還是執行大規模推論，TPU 都能提供可擴充的專用運算資源，並針對進階 AI 架構進行最佳化。

本文說明 TPU 在 Google Cloud中的角色、各個可用 TPU 版本的技術規格、效能特徵、定價考量事項，以及如何對應至 Compute Engine 機器系列。

什麼是 TPU？

TPU 是專為處理機器學習演算法核心的大型矩陣運算需求而設計的 ASIC。

TPU 系統的主要架構元件包括：

  • TensorCore：TPU 晶片的處理單元。每個 TensorCore 包含一或多個矩陣乘法單元 (MXU)，這些單元使用脈動陣列架構執行矩陣乘法，效率極高，此外還有純量和向量單元，用於控制流程和一般運算。

  • SparseCore：專門的資料流處理器，可加快稀疏運算。SparseCores 會與 TensorCores 搭配運作，有效處理大型嵌入資料表，因此非常適合加速建議模型和大規模嵌入作業。

  • 高頻寬記憶體 (HBM)：附加至 TPU 晶片的大型實體記憶體，提供大量記憶體頻寬。HBM 可訓練及服務較大的模型，並使用較大的批次大小。

  • TPU 執行個體：在 TPU 主機上執行的 Linux 運算執行個體，可直接存取底層 TPU 硬體，並提供高效能程式庫、SSH 連線能力和執行階段偵錯記錄檔。

  • TPU 配量和 Pod：TPU Pod 是實體連結 TPU 晶片的大型連續叢集。TPU 配量是 Pod 的邏輯分割區 (包含一或多個透過高速互連網路連線的主機)，用於執行單一工作負載。

  • XLA 編譯器：加速線性代數 (XLA) 編譯器會將機器學習計算圖編譯為最佳化的機器指令，在 TPU TensorCore 上執行。

TPU 版本和機器系列

Compute Engine 支援多個世代和版本的 TPU。下表列出各支援 TPU 版本的關鍵規格：

TPU 版本 機器系列 每顆晶片的尖峰運算效能 (TFLOP) TPU 記憶體和頻寬 每個晶片的核心數 每個晶片的互連 (ICI) 頻寬 每個晶片的網路頻寬 (DCN) 每個 Pod 的晶片數
TPU7x (Ironwood) tpu7x-standard

BF16：2,307

FP8：4,614

 192 GiB (7,380 GBps)

2 個 TensorCore

4 個 SparseCore

 1,200 GBps 100 Gbps 9,216
TPU v6e (Trillium) ct6e-standard

BF16：918

FP8：918

 32 GiB (1,638 GBps)

1 個 TensorCore

2 個 SparseCore

 800 GBps 100 Gbps 256
TPU v5p ct5p-hightpu

BF16：459

FP8：459

 95 GiB (2,765 GBps)

2 個 TensorCore

4 個 SparseCore

 1,200 GBps 50 Gbps 8,960

TPU 效能特徵

TPU 經過最佳化調整，可盡可能有效率地執行密集矩陣代數運算。如要選擇何時最適合使用 TPU 處理機器學習工作負載，請參考下列準則。

理想用途和工作負載需求

選擇合適的 TPU 模型取決於工作負載的運算特性、記憶體需求和擴充性需求。選取下列任一分頁標籤，即可查看建議。

AI/機器學習訓練

  • 預先訓練大型模型：從頭開始訓練大型模型。這些工作負載受運算和通訊限制。

    • TPU 必備功能：高尖峰 FLOPS、低精確度支援，以及高互連頻寬。
    • 建議使用的 TPU 版本：
      • TPU7x。這個版本可提供高 FP8 效能。
      • TPU v5p。這個版本支援大規模叢集設定。

  • 微調和精簡模型：使用參數效率方法調整現有模型，或訓練較小的模型變體。

    • TPU 需求：中等運算量和充足的記憶體容量。
    • 建議使用的 TPU 版本：
      • TPU v6e。這個版本適合用於標準訓練工作，可節省成本。
      • TPU7x。這個版本可大幅提升大規模微調的效能。

AI/機器學習推論

  • 線上提供模型：部署模型以進行即時互動，延遲時間至關重要。在連續產生權杖期間，這些工作負載會受到記憶體頻寬限制。

    • 必要 TPU 功能：高 HBM 頻寬和大型晶片內建 SRAM，可盡量縮短資料擷取延遲時間。
    • 建議使用的 TPU 版本：
      • TPU v6e。這個版本經過最佳化，可大規模提供服務，同時兼顧成本效益。
      • TPU7x。這個版本提供高 HBM 頻寬和大型 SRAM，可加速權杖生成。

  • 執行批次推論：離線處理大型資料集，主要目標是提高輸送量。在提示預先填入階段，這些工作負載會受到運算限制。

    • TPU 需求：高運算密度，盡可能提高每美元的處理量。
    • 建議使用的 TPU 版本：
      • TPU v6e。這個版本可有效節省成本，適合用於高輸送量的批次處理作業。
      • TPU7x。這個版本提供高運算密度，可快速處理大量批次作業。

推薦系統

  • 處理大型嵌入：訓練及放送使用大型嵌入資料表的推薦模型。這些工作負載受限於容量和通訊。
    • TPU 必備功能：專用的 SparseCore 硬體，可平行處理稀疏運算、大容量 HBM，以及支援卸載至主機記憶體。
    • 建議使用的 TPU 版本：
      • TPU7x。這個版本每個晶片包含四個 SparseCore，並提供最大的 HBM 容量。
      • TPU v6e。這個版本每顆晶片包含兩個 SparseCore，可提供經濟實惠的服務。
      • TPU v5p。這個版本每顆晶片包含四個 SparseCore，可進行大規模訓練。

強化學習

  • 執行強化學習迴圈：管理混合工作負載，這類工作負載需要緊密的迴圈，才能產生樣本及更新政策權重。

    • TPU 必備功能：低延遲 ICI，可在晶片間快速傳輸權重和啟用狀態，以及高頻寬主機連線。

    • 建議使用的 TPU 版本：

      這些 TPU 版本使用高速 3D 環面互連。

工作負載考量事項

TPU 可能不是下列工作負載的最佳選擇：

  • 需要頻繁邏輯分支或包含許多元素級別代數運算的線性代數程式。
  • 工作負載依附於 JAX 或 PyTorch 中的自訂運算子，且在 CPU 上執行。
  • 需要雙精度浮點 (FP64) 算術運算的工作負載。

費用考量事項

Google Cloud的 TPU 定價取決於下列因素：

  • TPU 版本

  • 位置

  • 計費模式

如需完整的定價詳細資料，請參閱「TPU 定價」頁面。

使用和購買模式

您可以根據可用性需求和成本容許度，使用下列任一消耗選項佈建 TPU 執行個體：

  • 隨選：標準即付即用 (PAYG) 計費，可立即執行。提供最大的生命週期彈性，但供應情形取決於實體資源限制。
  • Spot：以大幅折扣 (最高 70%) 價格提供多餘的 Google Cloud 容量。因為 Google Cloud 可以停止或刪除 Spot VM，並在 30 秒內取回容量，因此這些運算執行個體非常適合容錯和檢查點工作負載。
  • 彈性啟動： 可從專屬集區要求運算執行個體，最多七天，且可用性較高。
  • 未來預留項目：預留項目可確保容量，且與隨選價格相比，可享有折扣價。預留指定未來日期和時間的 TPU 容量，最長可達 90 天 (日曆模式)，或 1 年以上。如果Google Cloud 核准您的要求，您就能在預留期間使用容量。

承諾使用折扣 (CUD)

如果要求未來預留項目期限為 1 年以上，您可針對預留的 TPU 容量取得 CUD。將 CUD 套用至運算執行個體，即可在承諾使用期間內，享有比隨選費率更低的價格。

詳情請參閱「Compute Engine 適用的 CUD」。

後續步驟