GKE 中的 TPU 簡介

本文說明 Cloud TPU 如何與 Google Kubernetes Engine (GKE) 搭配運作，包括術語、Tensor Processing Unit (TPU) 的優點，以及工作負載排程考量。TPU 是 Google 開發的客製化特殊應用積體電路 (ASIC)，可加速使用 TensorFlow、PyTorch 和 JAX 等架構的機器學習工作負載。

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

閱讀本文前，請先確認您熟悉機器學習加速器的運作方式。詳情請參閱「Cloud TPU 簡介」。

在 GKE 中使用 TPU 的優點

GKE 全面支援 TPU 節點和節點集區生命週期管理，包括建立、設定及刪除 TPU VM。GKE 也支援 Spot VM，以及使用預留的 Cloud TPU。詳情請參閱「Cloud TPU 用量方案」。

在 GKE 中使用 TPU 的優點包括：

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

開始使用 Ironwood (TPU7x)

Ironwood (TPU7x) 是 Google 的第七代 TPU，專為大規模 AI 工作負載設計。如要進一步瞭解 Ironwood (TPU7x) 的優點，請參閱「關於 GKE 中的 Ironwood (TPU7x)」。

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 都啟動並執行後，Kubernetes Pod 就能排定在新 TPU 切片的 VM 上執行。

下圖顯示 v5litepod-16 (v5e) 多主機 TPU 配量。這個 TPU 切片有四個 VM。TPU 切片中的每個 VM 都有四個 TPU v5e 晶片，透過高速互連網路 (ICI) 連接，每個 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 slice 的範例，其中包含七部 v4-8 機器：

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

GKE 中 TPU 的特性

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

TPU 用量

為提高資源使用率和成本效益，同時兼顧工作負載效能，GKE 支援下列 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 晶片數量。舉例來說，您可以定義 TPU 晶片少於 64 個的小型拓撲，拓撲形式包括 2x2x2、2x2x4 或 2x4x4。如果使用超過 64 個 TPU 晶片的大型拓撲，指派給 {A}、{B} 和 {C} 的值必須符合下列條件：

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

機型命名

GKE 中 TPU 的機型命名方式，會因叢集模式和 TPU 版本而異：

  • GKE Standard：選取特定 Compute Engine 機型，例如 TPU Trillium (v6e) 的 ct6e-standard-1t。

  • GKE Autopilot：您不直接選取機型，而是使用工作負載資訊清單中的加速器類型要求 TPU。例如，TPU Trillium (v6e) 請使用 tpu-v6e-slice，TPU v5e 則使用 tpu-v5-lite-podslice。GKE Autopilot 接著會佈建基礎節點，並採用適當的機型來滿足要求。

如要瞭解各 TPU 版本適用的確切機器類型，請參閱「在 GKE 中規劃 TPU」中的表格。

特殊權限模式

如果您使用的 GKE 版本低於 1.28，則必須為容器設定特殊功能，才能存取 TPU。在 Standard 模式叢集中，您可以使用特殊權限模式授予這項存取權。特殊權限模式會覆寫 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 晶片。否則，Deployment 會失敗，因為 GKE 無法部分使用 TPU 資源。請參考下列情境：
    • 採用 2x4 拓撲的機型 ct5lp-hightpu-4t 包含兩個 TPU 配量節點，每個節點有四個 TPU 晶片，總共八個 TPU 晶片。使用這類機器時，您可以：
    • 無法在這個節點集區的節點上，部署需要八個 TPU 晶片的 GKE Pod。
    • 可以部署兩個 Pod，每個 Pod 需要四個 TPU 晶片，且每個 Pod 位於這個節點集區中的兩個節點之一。
    • 拓撲為 4x4 的 TPU v5e 在四個節點中含有 16 個 TPU 晶片。選取這項設定的 GKE Autopilot 工作負載，必須在每個副本中要求四個 TPU 晶片，副本數量為一到四個。
  • 在 Standard 叢集中，您可以在 VM 上排定多個 Kubernetes Pod，但每個 Pod 中只有一個容器可以存取 TPU 晶片。
  • 如要建立 kube-system Pod (例如 kube-dns)，每個 Standard 叢集都必須至少有一個非 TPU 切片節點集區。
  • 根據預設，TPU slice 節點具有 google.com/tpu taint，可防止非 TPU 工作負載排定在 TPU slice 節點上。不使用 TPU 的工作負載會在非 TPU 節點上執行，釋出 TPU 配量節點上的運算資源，供使用 TPU 的程式碼使用。請注意，taint 不保證 TPU 資源會充分利用。
  • GKE 會收集 TPU 切片節點上執行的容器所產生的記錄。詳情請參閱「Logging」。
  • Cloud Monitoring 提供 TPU 使用率指標，例如執行階段效能。詳情請參閱「觀測能力和指標」。
  • 您可以使用 GKE Sandbox 對 TPU 工作負載進行沙箱處理。GKE Sandbox 適用於 TPU 模型 v4 以上版本。詳情請參閱「GKE Sandbox」。

自動建立 TPU 節點集區

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

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

其他 Cloud TPU 類型則適用於所有 GKE 版本。如要進一步瞭解 TPU 適用的 GKE 版本，請參閱「驗證 GKE 中的 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 配量節點集區」。

收集排程

收集排程僅適用於 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，請參閱下列頁面：