關於 GKE 中的 TPU

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

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

開始使用 Ironwood (TPU7x)

Ironwood (TPU7x) 是 Google 的第七代 TPU，專為大規模 AI 工作負載設計。如要進一步瞭解 Ironwood (TPU7x)，請參閱下列資源：

下列每個叢集建立選項在叢集設定和工作負載排程方面，都提供不同程度的便利性和彈性：

  • 使用 Accelerated Processing Kit (XPK) 快速建立 GKE 叢集，並執行概念驗證和測試的工作負載。詳情請參閱 XPK 說明文件。
  • 使用 Google Cloud CLI 手動建立 GKE 叢集執行個體，以便精確自訂或擴充現有的正式版 GKE 環境

GKE 中與 TPU 相關的術語

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

  • Cloud TPU ICI 韌性：這項功能可提升光纖連結和光纖電路交換器 (OCS) 的容錯能力，這些連結和交換器會連接立方體之間的 TPU。詳情請參閱 TPU 架構。
  • TPU 立方體：互連 TPU 晶片的 4x4x4 拓撲。這僅適用於 3 元組 ({A}x{B}x{C}) 中的拓撲。
  • TPU 類型：Cloud TPU 類型，例如 v5e。
  • TPU 配量：位於同一 TPU Pod 內的一組晶片，透過高速晶片間互連 (ICI) 網路連線。配量會以晶片或 TensorCore 來說明，視 TPU 版本而定。
  • TPU 配量節點：由單一 VM 代表的 Kubernetes 節點，具有一或多個互連的 TPU 晶片。
  • TPU 節點集區：叢集內的一組 Kubernetes 節點，全部採用相同的 TPU 設定。
  • TPU 拓撲：TPU 配量中的 TPU 晶片數量和實體排列方式。
  • 不可分割：GKE 會將所有互連節點視為單一單位。在縮放作業期間，GKE 會將整組節點縮放為 0，並建立新節點。如果群組中的機器故障或終止，GKE 會以新單元的形式重新建立整組節點。
  • 不可變動：您無法手動將新節點新增至互連節點集。不過，您可以建立具有所需 TPU 拓撲的新節點集區，並在新節點集區上排定工作負載。

TPU 配量節點集區類型

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

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

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

多主機 TPU 配量節點集區

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

多主機 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 裝置，單一主機節點集區內的 VM 可以透過資料中心網路 (DCN) 通訊，但附加至 VM 的 TPU 不會互連。

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

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

GKE 中 TPU 的特性

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

TPU 用量

為提高資源用量和成本效益，同時兼顧工作負載效能，GKE 支援下列 TPU 用量選項：

如要選擇符合工作負載需求的用量選項，請參閱「關於 GKE 中 AI/機器學習工作負載的加速器用量選項」。

在 GKE 中使用 TPU 前，請先選擇最符合工作負載需求的消耗量選項。

拓撲

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

如果是排定在多主機 TPU 配量節點集區中的 TPU v4、v5p 和 Ironwood (TPU7x) (預先發布版)，您可以在 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。舉例來說，tpu7x-standard-4t 機型支援 Ironwood (TPU7x)，並包含四個 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 資源，部署作業會失敗。請參考下列情境：
    • 具有 2x4 拓撲的 ct5lp-hightpu-4t 機器類型包含兩個 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 自動建立節點集區

節點集區自動建立功能僅支援特定 GKE 版本中的下列 Cloud TPU：

  • TPU v3：1.31.0 以上版本。
  • TPU v5 和 TPU v4：1.29.0 以上版本。
  • TPU Trillium：1.31.1-gke.1146000 以上版本。
  • Ironwood (TPU7x) (搶先版)：1.34.1-gke.2541000 以上版本。

所有 GKE 版本都支援其他 Cloud TPU 類型。

Cloud TPU 節點集區自動調度資源

GKE 會透過下列其中一種方式，自動擴展或縮減使用叢集自動調整程式的 Cloud TPU 節點集區 (自動建立或手動建立)：

  • 單一主機 TPU 配量節點集區：GKE 會在現有節點集區中新增或移除 TPU 節點。節點集區可能包含任意數量的 TPU 節點，介於零和節點集區大小上限之間，而節點集區大小上限是由 --max-nodes 和 --total-max-nodes 自動調度資源標記決定。節點集區中的所有 TPU 節點都具有相同的機器類型和拓撲。如要進一步瞭解如何建立單一主機 TPU 配量節點集區，請參閱「建立單一主機 TPU 配量節點集區」。
  • 多主機 TPU 配量節點集區：GKE 會將節點集區從零原子式擴充至滿足 TPU 拓撲所需的節點數量。舉例來說，如果 TPU 節點集區的機器類型為 ct5lp-hightpu-4t，拓撲為 16x16，則節點集區一律會有 64 個節點或零個節點。如果節點集區中沒有 TPU 工作負載，GKE 就會縮減節點集區。如要縮減節點集區，GKE 會排除所有已排定的 Pod，並移除節點集區中的所有節點。如要進一步瞭解如何建立多主機 TPU 配量節點集區，請參閱建立多主機 TPU 配量節點集區。

工作負載政策

工作負載政策僅支援 Ironwood (TPU7x)。

工作負載政策是一種資源政策，可讓您設定基礎架構的實體位置。工作負載政策提供階層式設定選項，可進行更精細的控制。

Ironwood (TPU7x) 支援高處理量工作負載政策。這項政策可讓您執行下列操作：

  • 執行需要高處理量的分散式工作負載，例如高效能運算 (HPC) 或機器學習 (ML) 訓練。基礎架構已設定為將 TPU VM 置於同一處，以縮短 VM 之間的網路延遲時間。
  • 定義 TPU VM 的維護策略，盡量減少工作負載中斷。

收藏排程

集合排程僅支援 TPU Trillium。

在 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，請參閱下列頁面：