關於 GKE 中的 TPU

本頁說明 Cloud TPU 如何與 Google Kubernetes Engine (GKE) 搭配運作，包括術語、張量處理單元 (TPU) 的優點，以及工作負載排程考量。TPU 是 Google 專門打造的特殊應用積體電路 (ASIC)，用於加速處理使用 TensorFlow、PyTorch 和 JAX 等架構的機器學習工作負載。

本頁面適用於平台管理員和營運人員，以及執行機器學習 (ML) 模型的資料和 AI 專家，這些模型具有大規模、長時間執行或以矩陣運算為主等特徵。如要進一步瞭解內容中提及的常見角色和範例工作，請參閱「常見的 GKE Enterprise 使用者角色和工作」。 Google Cloud

閱讀本頁內容前，請務必先熟悉機器學習加速器的運作方式。詳情請參閱「Cloud TPU 簡介」。

在 GKE 中使用 TPU 的好處

GKE 全面支援 TPU 節點和節點集區生命週期管理，包括建立、設定及刪除 TPU VM。GKE 也支援現成 VM，以及使用預留的 Cloud TPU。詳情請參閱「Cloud TPU 消耗選項」。 在 GKE 中使用 TPU 的好處包括：

  • 一致的作業環境：您可以使用單一平台處理所有機器學習和其他工作負載。
  • 自動升級：GKE 會自動更新版本，減少營運負擔。
  • 負載平衡：GKE 會分配負載，藉此縮短延遲時間並提升可靠性。
  • 彈性調度：GKE 會自動調度 TPU 資源，滿足工作負載需求。
  • 資源管理：透過 Kubernetes 原生工作佇列系統 Kueue，您可以使用佇列、搶占、優先順序和公平分享，管理機構內多個租戶的資源。
  • 沙箱選項：GKE Sandbox 可透過 gVisor 保護工作負載。詳情請參閱 GKE Sandbox。

使用 TPU Trillium 的好處

Trillium 是 Google 的第六代 TPU，Trillium 具備下列優勢：

  • 與 TPU v5e 相比，Trillium 的每個晶片運算效能更高。
  • Trillium 增加了高頻寬記憶體 (HBM) 容量和頻寬，與 TPU v5e 相比，晶片間互連 (ICI) 頻寬也更高。
  • Trillium 配備第三代 SparseCore，這是專門的加速器，可處理進階排名和推薦工作負載中常見的超大型嵌入。
  • Trillium 的能源效率比 TPU v5e 高出 67% 以上。
  • 在單一高頻寬、低延遲的 TPU 區塊中，Trillium 最多可擴充至 256 個 TPU。
  • Trillium 支援集合排程。集合排程可讓您宣告一組 TPU (單一主機和多主機 TPU 節點集區)，確保推論工作負載需求的高可用性。

在所有技術介面 (例如 API 和記錄) 中，以及 GKE 說明文件的特定部分，我們使用 v6e 或 TPU Trillium (v6e) 來指稱 Trillium TPU。如要進一步瞭解 Trillium 的優點，請參閱這篇 Trillium 公告網誌文章。如要開始設定 TPU，請參閱「在 GKE 中規劃 TPU」。

GKE 中與 TPU 相關的術語

本頁面會使用下列與 TPU 相關的術語：

  • TPU 類型：Cloud TPU 類型，例如 v5e。
  • TPU 配量節點：由單一 VM 代表的 Kubernetes 節點，其中包含一或多個互連的 TPU 晶片。
  • TPU 節點集區：叢集中具有相同 TPU 設定的一組 Kubernetes 節點。
  • TPU 拓撲：TPU 配量中 TPU 晶片的數量和實體排列方式。
  • 不可分割：GKE 會將所有互連節點視為單一單位。在縮放作業期間，GKE 會將整組節點縮放為 0，並建立新節點。如果群組中的機器故障或終止，GKE 會以新單元的形式重新建立整組節點。
  • Immutable：您無法手動將新節點新增至互連節點集。不過，您可以建立具有所需 TPU 拓撲的新節點集區，並在新節點集區上排定工作負載。

TPU 配量節點集區類型

GKE 支援兩種類型的 TPU 節點集區：

TPU 類型和拓撲會決定 TPU 配量節點是多主機還是單一主機。建議您採取以下做法：

  • 對於大規模模型，請使用多主機 TPU 配量節點
  • 如果是小型模型，請使用單一主機 TPU 配量節點
  • 如要進行大規模訓練或推論，請使用 Pathways。Pathways 可讓單一 JAX 用戶端協調多個大型 TPU 配量的工作負載，簡化大規模機器學習運算。詳情請參閱「路徑」。

多主機 TPU 配量節點集區

多主機 TPU 配量節點集區是包含兩個以上互連 TPU VM 的節點集區。每個 VM 都會連接一個 TPU 裝置，多主機 TPU 配量中的 TPU 會透過高速互連網路 (ICI) 連線。建立多主機 TPU 節點集區後，您就無法在其中新增節點。舉例來說，您無法建立 v4-32 節點集區，然後在節點集區中新增其他 Kubernetes 節點 (TPU VM)。如要在 GKE 叢集中新增 TPU 節點，請務必建立新的節點集區。

多主機 TPU 配量節點集區中的 VM 會視為單一不可分割的單位。如果 GKE 無法在配量中部署一個節點，系統就不會部署 TPU 配量節點中的任何節點。

如果多主機 TPU 配量中的節點需要修復，GKE 會關閉 TPU 配量中的所有 VM，強制驅逐工作負載中的所有 Kubernetes Pod。TPU 配量中的所有 VM 啟動並執行後，即可在新的 TPU 配量中，將 Kubernetes Pod 排定至 VM。

下圖顯示 v5litepod-16 (v5e) 多主機 TPU 配量。這個 TPU 配量有四個 VM。TPU 配量中的每個 VM 都有四個透過高速互連網路 (ICI) 連線的 TPU v5e 晶片，每個 TPU v5e 晶片都有一個 TensorCore。

多主機 TPU 配量圖

下圖顯示的 GKE 叢集包含一個 TPU v5litepod-16 (v5e) TPU 節點 (拓撲：4x4) 和一個 TPU v5litepod-8 (v5e) 節點 (拓撲：2x4)：

TPU v5e Pod 架構圖

單一主機 TPU 配量節點集區

單一主機配量節點集區包含一或多個獨立的 TPU VM。每個 VM 都會連接一個 TPU 裝置，單一主機 Slice 節點集區中的 VM 可以透過資料中心網路 (DCN) 通訊，但附加至 VM 的 TPU 不會互連。

下圖顯示單一主機 TPU 節點的範例，其中包含七部 v4-8 機器：

單一主機切片節點集區圖表

GKE 中 TPU 的特性

TPU 具有獨特特性，因此需要特別規劃和設定。

拓撲

拓撲會定義 TPU 切片中 TPU 的實體排列方式。 GKE 會根據 TPU 版本，在二維或三維拓撲中佈建 TPU 配量。您會指定每個維度的 TPU 晶片數量做為拓撲，如下所示：

如要為多主機 TPU 配量節點集區中排定的 TPU v4 和 v5p 定義拓撲，請使用 3 元組 ({A}x{B}x{C})，例如 4x4x4。{A}x{B}x{C} 的乘積會定義節點集區中的 TPU 晶片數量。舉例來說，您可以定義少於 64 個 TPU 晶片的拓撲，並使用 2x2x2、2x2x4 或 2x4x4 等拓撲形式。如果您使用的拓撲大於 64 個 TPU 晶片，則指派給 {A}、{B} 和 {C} 的值必須符合下列條件：

  • {A}、{B} 和 {C} 必須是四的倍數。
  • v4 支援的最大拓撲為 12x16x16，v5p 則為 16x16x24。
  • 指派的值必須符合 A ≤ B ≤ C 模式。例如 4x4x8 或 8x8x8。

機型

支援 TPU 資源的機器類型會遵循命名慣例，其中包含 TPU 版本和每個節點配量的 TPU 晶片數量，例如 ct<version>-hightpu-<node-chip-count>t。舉例來說，ct5lp-hightpu-1t 機器類型支援 TPU v5e，且只包含一個 TPU 晶片。

特殊權限模式

如果您使用 1.28 之前的 GKE 版本，必須為容器設定特殊功能，才能存取 TPU。在標準模式叢集中，您可以使用特權模式授予這項存取權。特殊權限模式會覆寫 securityContext 中的許多其他安全性設定。詳情請參閱在不使用特權模式的情況下執行容器。

1.28 以上版本不需要特殊權限模式或特殊功能。

GKE 中的 TPU 運作方式

Kubernetes 資源管理和優先順序設定會將 TPU 上的 VM 視為其他 VM 類型。如要要求 TPU 晶片，請使用資源名稱 google.com/tpu：

    resources:
        requests:
          google.com/tpu: 4
        limits:
          google.com/tpu: 4

在 GKE 中使用 TPU 時，請考量下列 TPU 特性：

  • VM 最多可存取 8 個 TPU 晶片。
  • TPU 配量包含固定數量的 TPU 晶片，數量取決於您選擇的 TPU 機器類型。
  • 要求的 google.com/tpu 數量必須等於 TPU 配量節點上可用的 TPU 晶片總數。GKE Pod 中要求 TPU 的任何容器都必須使用節點中的所有 TPU 晶片。否則，由於 GKE 無法部分耗用 TPU 資源，部署作業會失敗。請參考下列情境：
    • 機器類型 ct5lp-hightpu-4t 搭配 2x4 拓撲包含兩個 TPU 配量節點，每個節點有四個 TPU 晶片，總共八個 TPU 晶片。使用這類機器類型，您可以：
      • 無法在這個節點集區的節點上，部署需要八個 TPU 晶片的 GKE Pod。
      • 可部署兩個 Pod，每個 Pod 各需四個 TPU 晶片，且每個 Pod 位於這個節點集區的其中一個節點上。
    • 拓撲為 4x4 的 TPU v5e 在四個節點中共有 16 個 TPU 晶片。選取這項設定的 GKE Autopilot 工作負載，必須在每個副本中要求四個 TPU 晶片，副本數量為一到四個。
  • 在標準叢集中，多個 Kubernetes Pod 可排定在 VM 上執行，但每個 Pod 中只有一個容器可以存取 TPU 晶片。
  • 如要建立 kube-system Pod (例如 kube-dns)，每個標準叢集都必須至少有一個非 TPU 節點集區。
  • 根據預設，TPU 節點具有 google.com/tpu taint ，可防止非 TPU 工作負載排定在 TPU 節點上執行。不使用 TPU 的工作負載會在非 TPU 節點上執行，將 TPU 節點上的運算資源空出來，供使用 TPU 的程式碼使用。請注意，汙點無法保證 TPU 資源得到充分運用。
  • GKE 會收集在 TPU 節點上執行的容器所發出的記錄。詳情請參閱「記錄」。
  • 您可以在 Cloud Monitoring 中查看 TPU 使用率指標，例如執行階段效能。詳情請參閱「可觀測性和指標」。
  • 您可以使用 GKE Sandbox 將 TPU 工作負載放入沙箱。GKE Sandbox 適用於 TPU 模型 v4 以上版本。詳情請參閱「GKE Sandbox」。

如何排定收藏內容的發布時間

在 TPU Trillium 中，您可以使用集合排程，將 TPU 節點切片分組。將這些 TPU 節點分組後，就能更輕鬆地調整副本數量，以滿足工作負載需求。 Google Cloud 會控管軟體更新，確保集合內有足夠的切片，隨時可處理流量。

TPU Trillium 支援單一主機和多主機節點集區的集合排程，可執行推論工作負載。以下說明收集排程行為如何取決於所用 TPU 配量類型：

  • 多主機 TPU 配量：GKE 會將多主機 TPU 配量分組，形成集合。每個 GKE 節點集區都是這個集合中的副本。如要定義集合，請建立多主機 TPU 配量，並為集合指派專屬名稱。如要將更多 TPU 節點加入集合，請建立另一個多主機 TPU 節點集區，並使用相同的集合名稱和工作負載類型。
  • 單一主機 TPU 配量：GKE 會將整個單一主機 TPU 配量節點集區視為一個集合。如要將更多 TPU 配量新增至集合，可以調整單一主機 TPU 配量節點集區的大小。

集合排程有下列限制：

  • 您只能為 TPU Trillium 安排收集作業。
  • 您只能在建立節點集區時定義集合。
  • 不支援 Spot VM。
  • 包含多主機 TPU 配量節點集區的集合，必須為集合內的所有節點集區使用相同的機器類型、拓撲和版本。

您可以在下列情況設定集合排程：

後續步驟

如要瞭解如何在 GKE 中設定 Cloud TPU，請參閱下列頁面：