在 GKE Standard 部署 TPU 工作負載

本文是基礎文件,說明如何使用 Google Kubernetes Engine (GKE) 中的 TPU,加速機器學習 (ML) 工作負載。TPU 專為矩陣乘積處理而設計,例如大規模深度學習模型訓練。TPU 經過最佳化,可處理機器學習的大量資料集和複雜模型,因此效能更出色,機器學習工作負載的成本效益和能源效率也更高。本指南將說明如何使用 Cloud TPU 加速器部署 ML 工作負載、設定 TPU 配額、設定執行 TPU 的節點集區升級作業,以及監控 TPU 工作負載指標。

本教學課程適用於機器學習 (ML) 工程師、平台管理員和營運人員,他們有興趣使用 Kubernetes 容器協調程序,管理使用 TPU 的大規模模型訓練、調整和推論工作負載。如要進一步瞭解 Google Cloud 內容中提及的常見角色和範例工作,請參閱「常見的 GKE 使用者角色和工作」。

閱讀本頁內容前,請務必先熟悉下列事項:

事前準備

開始之前,請務必先完成下列工作:

  • 啟用 Google Kubernetes Engine API。
  • 啟用 Google Kubernetes Engine API
  • 如要使用 Google Cloud CLI 執行這項工作,請安裝並初始化 gcloud CLI。如果您先前已安裝 gcloud CLI,請執行 gcloud components update 指令,取得最新版本。較舊的 gcloud CLI 版本可能不支援執行本文件中的指令。

規劃 TPU 設定

根據模型和所需記憶體容量規劃 TPU 設定。如要使用本指南在 TPU 上部署工作負載,請先完成「規劃 TPU 設定」中的規劃步驟。

確認您有 TPU 配額

下列各節說明如何在 GKE 中使用 TPU 時,確保配額充足。

以量計價或 Spot VM 的配額

如要使用隨選或 Spot VM 建立 TPU 配量節點集區,您必須在要使用的區域中,擁有充足的 TPU 配額。

建立會耗用 TPU 預留項目的 TPU 配量節點集區時,不需要任何 TPU 配額。1 預留 TPU 時,您可以安全地略過這個部分。

在 GKE 中建立隨選或 Spot TPU 配量節點集區時,需要 Compute Engine API 配額。Compute Engine API 配額 (compute.googleapis.com) 與 Cloud TPU API 配額 (tpu.googleapis.com) 不同,後者是在使用 Cloud TPU API 建立 TPU 時的必要配額。

如要查看 TPU 的 Compute Engine API 配額限制和目前用量,請按照下列步驟操作:

  1. 前往 Google Cloud 控制台的「Quotas」(配額) 頁面:

    前往配額頁面

  2. 在「篩選器」方塊中,執行下列操作:

    1. 請根據 TPU 版本和機型,使用下表選取並複製配額的屬性。舉例來說,如果您打算建立機器類型開頭為 ct5lp- 的隨選 TPU v5e 節點,請輸入 Name: TPU v5 Lite PodSlice chips。

      TPU 版本、 機型開頭為 隨選執行個體的配額屬性和名稱 Spot 2 執行個體的配額屬性和名稱
      TPU v3、
      ct3-
      Dimensions (e.g. location):
      tpu_family:CT3
      不適用
      TPU v3、
      ct3p-
      Dimensions (e.g. location):
      tpu_family:CT3P
      不適用
      TPU v4、
      ct4p-
      Name:
      TPU v4 PodSlice chips
      Name:
      Preemptible TPU v4 PodSlice chips
      TPU v5e、
      ct5lp-
      Name:
      TPU v5 Lite PodSlice chips
      Name:
      Preemptible TPU v5 Lite Podslice
      chips
      TPU v5p、
      ct5p-
      Name:
      TPU v5p chips
      Name:
      Preemptible TPU v5p chips
      TPU Trillium、
      ct6e-
      Dimensions (e.g. location):
      tpu_family:CT6E
      Name:
      Preemptible TPU slices v6e
      Ironwood (TPU7x)、
      tpu7x-standard-4t
      Dimensions (e.g. location):
      tpu_family:tpu7x
      Name:
      Preemptible TPU slices tpu7x
    2. 選取「Dimensions (e.g. locations)」(維度,例如位置) 屬性,然後輸入 region:,接著輸入您要在 GKE 中建立 TPU 的區域名稱。舉例來說,如果您打算在 us-west4-a 可用區建立 TPU 節點,請輸入 region:us-west4。TPU 配額屬於區域配額,因此同一區域內的所有可用區會共用配額。

如果沒有任何配額符合您輸入的篩選條件,表示專案尚未獲得您所需區域的任何指定配額,因此您必須申請調整 TPU 配額。

建立 TPU 預留項目時,對應配額的上限和目前用量值都會增加,增幅為 TPU 預留項目中的晶片數量。舉例來說,如果為 16 個 TPU v5e 晶片建立預訂,且這些晶片的機器類型開頭為 ct5lp-,則相關區域中 TPU v5 Lite PodSlice chips 配額的「限制」和「目前用量」都會增加 16。

  1. 建立 TPU 配量節點集區時,請使用 --reservation 和 --reservation-affinity=specific 旗標建立預留執行個體。購買承諾時,即可使用 TPU 預留項目。↩

  2. 建立 TPU 節點集區時,請使用 --spot 標記建立 Spot 執行個體。↩

其他 GKE 資源的配額

您可能需要在 GKE 建立資源的區域中,提高下列 GKE 相關配額。

  • Persistent Disk SSD (GB) 配額:每個 Kubernetes 節點的開機磁碟預設需要 100 GB。因此,這項配額至少應設為您預期建立的 GKE 節點數量上限乘以 100 GB (節點 * 100 GB)。
  • 使用中的 IP 位址配額:每個 Kubernetes 節點會耗用一個 IP 位址。 因此,這項配額至少應設為您預期建立的 GKE 節點數量上限。
  • 確認 max-pods-per-node 與子網路範圍一致:每個 Kubernetes 節點都會使用 Pod 的次要 IP 範圍。舉例來說,max-pods-per-node32 個節點需要 64 個 IP 位址,也就是每個節點的 /26 子網路。請注意,這個範圍不應與任何其他叢集共用。為避免 IP 位址範圍用盡,請使用 --max-pods-per-node 標記,限制可在節點上排程的 Pod 數量。max-pods-per-node 的配額應至少設為您預期建立的 GKE 節點數量上限。

如要申請提高配額,請參閱「要求調整配額」。

確認預訂是否可用

如要使用預留項目建立 TPU 配量節點集區,預留項目必須在建立節點集區時,有足夠的可用 TPU 晶片。

如要查看專案中的預留項目,以及 TPU 預留項目中有多少 TPU 晶片可用,請查看預留項目清單。

建立叢集

gcloud

在有可用 TPU 的區域中,以標準模式建立 GKE 叢集。

最佳做法:

使用區域叢集,確保 Kubernetes 控制層具備高可用性。

gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --cluster-version VERSION

更改下列內容:

  • CLUSTER_NAME:新叢集的名稱。
  • LOCATION:提供 TPU 容量的區域。
  • VERSION:GKE 版本,必須支援您要使用的機型。請注意,預設 GKE 版本可能不支援目標 TPU。 如要瞭解各 TPU 機型適用的最低 GKE 版本,請參閱 GKE 中的 TPU 可用性。

Cluster Toolkit

  1. 設定 Cluster Toolkit。
  2. 使用 Cluster Toolkit 藍圖 (例如 gke-tpu-v6e 或 gke-tpu-7x;如需 TPU v4、v5e 和 v5p 藍圖,請參閱 Cluster Toolkit GitHub 存放區) 部署 GKE TPU 叢集:

    gcluster deploy examples/gke-tpu-7x/gke-tpu-7x.yaml \
      --vars project_id=PROJECT_ID,deployment_name=DEPLOYMENT_NAME,zone=ZONE
    

    更改下列內容:

    • DEPLOYMENT_NAME:Cluster Toolkit 部署作業的名稱 (也是 GKE 叢集的名稱)。
    • PROJECT_ID:您的 Google Cloud 專案 ID。
    • ZONE:根據您要使用的 TPU 版本,選擇適用的可用區名稱。如要找出可用位置,請參閱「GKE 中的 TPU 可用性」。如要使用預留容量,請務必使用預留容量的可用區。

您已建立 GKE 叢集,且節點集區包含指定的 TPU 資源。現在可以按照本文「在 TPU 節點上執行工作負載」一節中的步驟操作。

佈建 TPU

如要在 GKE 中佈建 TPU,您可以選擇下列設定選項:
  • 手動建立節點集區:您可以建立具有特定 TPU 版本和拓撲的節點集區。
  • 使用 GKE 自動佈建節點功能:您可以在叢集層級啟用自動佈建節點功能,然後在 Pod 的資訊清單中,使用 nodeSelector 指定 TPU 版本和拓撲。當待處理的 Pod 符合這些選取器時,GKE 會自動建立符合要求的節點集區。這個方法需要為 TPU 設定叢集層級的資源限制。
  • 定義自訂 ComputeClass:您可以使用自訂 ComputeClass 要求 TPU。平台管理員可透過自訂 ComputeClass 定義節點設定的階層,供 GKE 在決定節點資源調度時優先採用,確保工作負載在您選取的硬體上執行。

手動建立節點集區

您可以建立單一或多主機 TPU 配量節點集區。

建立單一主機 TPU 配量節點集區

您可以使用 Google Cloud CLI、Terraform 或 Google Cloud 控制台,建立單一主機 TPU 節點集區。

gcloud

gcloud container node-pools create NODE_POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    [--sandbox=type=gvisor]

更改下列內容:

  • NODE_POOL_NAME:新節點集區的名稱。
  • LOCATION:根據您要使用的 TPU 版本,選擇適用的可用區名稱。如要找出可用位置,請參閱「GKE 中的 TPU 可用性」。
  • CLUSTER_NAME:叢集名稱。
  • NODE_ZONES:以半形逗號分隔的一或多個區域清單,GKE 會在這些區域中建立節點集區。

    選用:使用 AI 區域,例如 us-central1-ai1a。AI 可用區是專門的地點,經過最佳化調整,可處理 Google Cloud 區域內的 AI/機器學習工作負載。

  • MACHINE_TYPE:TPU 版本和類型。舉例來說,請使用 tpu7x-standard-4t 代表 Ironwood (TPU7x)。

您也可以視需要使用下列旗標:

  • --num-nodes=NUM_NODES:每個可用區中節點集區的初始節點數量。

    最佳做法:

    如果節點集區使用 enable-autoscaling 旗標,請將 num-nodes 設為 0,以便自動調度器在工作負載需要時,立即佈建額外節點。

  • --reservation=RESERVATION_NAME:GKE 在建立節點集區時使用的預留項目名稱。如果省略這個旗標,GKE 會使用可用的 TPU。如要進一步瞭解 TPU 預留資源,請參閱「關於 Cloud TPU 預留資源」。

  • --node-labels cloud.google.com/gke-workload-type=HIGH_AVAILABILITY:告知 GKE 單一主機 TPU 配量節點集區是集合的一部分。如果符合下列條件,請使用這個標記:

    • 節點集區會在新的節點集區中執行推論工作負載。
    • 節點集區使用 TPU Trillium。
    • 節點集區未使用 Spot VM。

    如要進一步瞭解如何管理集合排程,請參閱「在單一主機 TPU 切片中管理集合排程」。

  • --enable-autoscaling:建立已啟用自動調度資源的節點集區。 需要下列額外旗標:

    • --total-min-nodes=TOTAL_MIN_NODES:節點集區中所有節點的數量下限。
    • --total-max-nodes=TOTAL_MAX_NODES:節點集區中所有節點的數量上限。
    • --location-policy=ANY:優先使用未使用的預留項目,並降低 Spot VM 的先占風險。
  • --spot:將節點集區設為使用Spot VM,用於節點集區中的節點。節點集區建立後即無法變更這項設定。

  • --flex-start:將節點集區設為使用彈性啟動型 VM。使用「彈性啟動」消耗選項建立彈性啟動 VM。詳情請參閱「使用 TPU 和彈性啟動 VM 執行小型批次工作負載」。

如需可指定的完整旗標清單,請參閱gcloud container clusters create 參考資料。

Terraform

  1. 請確認您使用的是 4.84.0 以上版本的 google 供應商。
  2. 在 Terraform 設定中新增下列區塊:
resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]

  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
    flex_start = false
  }
}

更改下列內容:

  • NODE_POOL_RESOURCE_NAME:Terraform 範本中的節點集區資源名稱。
  • PROJECT_ID:您的專案 ID。
  • CLUSTER_NAME:現有叢集的名稱。
  • POOL_NAME:要建立的節點集區名稱。
  • CLUSTER_LOCATION:叢集的運算可用區。指定可使用 TPU 版本的區域。詳情請參閱「選取 TPU 版本和拓撲」。
  • NODE_ZONES:以半形逗號分隔的一或多個區域清單,GKE 會在這些區域中建立節點集區。

    選用:使用 AI 區域,例如 us-central1-ai1a。AI 可用區是專門的地點,經過最佳化調整,可處理 Google Cloud 區域內的 AI/機器學習工作負載。

  • MACHINE_TYPE:要使用的 TPU 機型類型。如要查看與 TPU 相容的機器類型,請參閱「選擇 TPU 版本」一文中的表格。

您也可以視需要使用下列變數:

  • autoscaling:建立已啟用自動調度資源的節點集區。如果是單一主機 TPU 節點,GKE 會在 TOTAL_MIN_NODES 和 TOTAL_MAX_NODES 值之間進行調整。
    • TOTAL_MIN_NODES:節點集區中所有節點的數量下限。除非也指定自動調度,否則這個欄位為選填。
    • TOTAL_MAX_NODES:節點集區中所有節點的數量上限。除非也指定自動調度,否則這個欄位為選填。
  • RESERVATION_NAME:如果您使用「關於 Cloud TPU 預訂」,這是建立節點集區時要使用的預訂資源標籤清單。如要進一步瞭解如何填寫 RESERVATION_LABEL_VALUES 中的 reservation_affinity 欄位,請參閱 Terraform 供應商。
  • spot:將節點集區設為使用 Spot VM 做為 TPU 節點。節點集區建立後即無法變更這項設定。詳情請參閱「Spot VM」。
  • flex_start:將節點集區設為使用彈性啟動用量方案。如果已啟用 spot,則無法設為 true。GKE 1.33.0-gke.1712000 以上版本支援彈性啟動。

控制台

如要建立含 TPU 的節點集區,請按照下列步驟操作:

  1. 前往 Google Cloud 控制台的「Google Kubernetes Engine」頁面。

    前往「Google Kubernetes Engine」

  2. 在叢集清單中,按一下您要修改的叢集名稱。

  3. 按一下「Nodes」(節點) 分頁標籤。

  4. 按一下「建立由使用者管理的節點集區」。

  5. 在「節點集區詳細資料」部分中,勾選「指定節點位置」方塊。

  6. 根據要使用的 TPU 版本選取可用區。如要找出可用區,請參閱「GKE 中的 TPU 可用性」。

  7. 按一下導覽窗格中的「Nodes」(節點)。

  8. 在「機器設定」部分,選取「TPU」。

  9. 在「系列」下拉式選單中,選取下列任一選項:

    • CT3:TPU v3,單一主機裝置
    • CT3P:TPU v3,多主機 Pod 配量
    • CT4P:TPU v4
    • CT5LP:TPU v5e
    • CT5P:TPU v5p
    • CT6E:TPU Trillium (v6e)
  10. 在「機型」下拉式選單中,選取要用於節點的機器名稱。請參閱「選擇 TPU 版本」表格,瞭解如何定義機型和 TPU 拓撲,以建立單一主機 TPU 配量節點集區。

  11. 在「TPU 拓撲」下拉式選單中,選取 TPU 配量的實體拓撲。

  12. 在「需要變更」對話方塊中,按一下「進行變更」。

  13. 確認「開機磁碟類型」為「標準永久磁碟」或「SSD 永久磁碟」。

  14. 或者,您也可以勾選「啟用 Spot VM 上的節點」核取方塊,為節點集區中的節點使用 Spot VM。

  15. 點選「建立」。

建立多主機 TPU 配量節點集區

建立多主機 TPU 配量節點集區的步驟,取決於您使用的是 Ironwood (TPU7x) 還是較早的 TPU 版本。

Ironwood (TPU7x)

您可以使用 Google Cloud CLI 或 Terraform,在 Ironwood (TPU7x) 版本中建立多主機 TPU 配量節點集區:

gcloud

如要使用 Ironwood (TPU7x) 建立多主機 TPU 配量節點集區,請先建立工作負載政策。

  1. 建立工作負載政策:

    gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
        --type=HIGH_THROUGHPUT \
        --accelerator-topology=TPU_TOPOLOGY \
        --project=PROJECT_ID \
        --region=REGION
    

    更改下列內容:

    • WORKLOAD_POLICY_NAME:工作負載政策的名稱。
    • TPU_TOPOLOGY:TPU Ironwood (TPU7x) 拓撲。 例如,2x2x2。如要查看所有支援的 Ironwood (TPU7x) 拓撲,請參閱拓撲部分。
    • PROJECT_ID:您的 Google Cloud 專案 ID。
    • REGION:工作負載政策的區域。工作負載政策是區域性資源,可供共用相同拓撲的節點集區重複使用。
  2. 使用工作負載政策建立節點集區:

    gcloud container node-pools create NODE_POOL_NAME \
      --cluster=CLUSTER_NAME \
      --machine-type=tpu7x-standard-4t \
      --placement-policy=WORKLOAD_POLICY_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --node-locations=NODE_ZONE \
      --project=PROJECT_ID \
      --reservation=RESERVATION_NAME \
      --reservation-affinity=specific
    

    更改下列內容:

    • NODE_POOL_NAME:新節點集區的名稱。
    • CLUSTER_NAME:GKE 叢集的名稱。
    • WORKLOAD_POLICY_NAME:您建立的工作負載政策名稱。
    • CONTROL_PLANE_LOCATION:叢集控制層的 Compute Engine 位置。如果是區域叢集,請提供區域;如果是可用區叢集,請提供可用區。
    • NODE_ZONE:根據您要使用的 TPU 版本,選擇適用的可用區名稱。

      選用:使用 AI 區域,例如 us-central1-ai1a。AI 可用區是專門的地點,經過最佳化調整,可處理 Google Cloud 區域內的 AI/機器學習工作負載。

      如要找出可用位置,請參閱「GKE 中的 TPU 可用性」。
    • PROJECT_ID:您的 Google Cloud 專案 ID。
    • RESERVATION_NAME:要使用的預訂名稱。

    在這個指令中,--tpu-topology 旗標已由 --placement-policy 旗標取代。

Terraform

  1. 請確認您使用的是 4.84.0 以上版本的 google 供應商。
  2. 建立工作負載政策:

    resource "google_compute_resource_policy" {
      name   = "WORKLOAD_POLICY_NAME"
      region = CLUSTER_LOCATION
      workload_policy {
        type = "HIGH_THROUGHPUT"
        accelerator_topology = "TPU_TOPOLOGY"
      }
    }
    

    更改下列內容:

    • WORKLOAD_POLICY_NAME:工作負載政策的名稱。
    • CLUSTER_LOCATION:叢集的運算位置。建議您使用區域性叢集,提高 Kubernetes 控制層的可靠性。您也可以使用區域叢集。 詳情請參閱「選取 TPU 版本和拓撲」。
    • TPU_TOPOLOGY:TPU Ironwood (TPU7x) 拓撲。 例如,2x2x2。如要查看所有支援的 Ironwood (TPU7x) 拓撲,請參閱「規劃 TPU」。

    如要進一步瞭解 google_compute_resource_policy 參考資料,請參閱 Terraform Provider。

  3. 在 Terraform 設定中,新增下列區塊:

    resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
      provider           = google
      project            = PROJECT_ID
      cluster            = CLUSTER_NAME
      name               = POOL_NAME
      location           = CLUSTER_LOCATION
      node_locations     = [NODE_ZONES]
      initial_node_count = NUM_NODES
    
      autoscaling {
        max_node_count = MAX_NODES
        location_policy      = "ANY"
      }
      node_config {
        machine_type = MACHINE_TYPE
        reservation_affinity {
          consume_reservation_type = "SPECIFIC_RESERVATION"
          key = "compute.googleapis.com/reservation-name"
          values = [RESERVATION_LABEL_VALUES]
        }
        flex_start = false
        spot = true
      }
    
      placement_policy {
        policy_name = WORKLOAD_POLICY_NAME
      }
    }
    

    更改下列內容:

    • NODE_POOL_RESOURCE_NAME:Terraform 範本中的節點集區資源名稱。
    • PROJECT_ID:您的專案 ID。
    • CLUSTER_NAME:要新增節點集區的現有叢集名稱。
    • POOL_NAME:要建立的節點集區名稱。
    • NODE_ZONES:以半形逗號分隔的一或多個區域清單,GKE 會在這些區域中建立節點集區。

      選用:使用 AI 區域,例如 us-central1-ai1a。AI 可用區是專門的地點,經過最佳化調整,可處理 Google Cloud 區域內的 AI/機器學習工作負載。

    • NUM_NODES:節點集區中的節點數量。這個值必須是零,或是 TPU 晶片數量除以四的結果,因為在多主機 TPU 配量中,每個 TPU 配量節點都有四個晶片。舉例來說,如果 TPU_TOPOLOGY 是 4x8,則有 32 個晶片,因此 NUM_NODES 必須是 8。如要進一步瞭解 TPU 拓撲,請參閱「選擇 TPU 版本」一文中的表格。
    • TPU_TOPOLOGY:這表示 TPU 配量選取的實體拓撲。拓撲格式取決於您使用的 TPU 版本。如要進一步瞭解 TPU 拓撲,請參閱「選擇拓撲」一文中的表格。

    您也可以視需要使用下列變數:

    • RESERVATION_NAME:如果您使用 TPU 預訂,請提供建立節點集區時要使用的預訂資源標籤清單。如要進一步瞭解如何填入「RESERVATION_LABEL_VALUES」欄位中的「reservation_affinity」,請參閱 Terraform 供應商。
    • autoscaling:建立已啟用自動調度資源的節點集區。當 GKE 擴充多主機 TPU 配量節點集區時,會不可分割地將節點集區從零擴充至最大大小。
      • MAX_NODES:節點集區的大小上限。這個值必須等於 TPU_TOPOLOGY ({A}x{B}x{C}) 中定義的值的乘積,除以每個 VM 中的晶片數量。舉例來說,如果 TPU_TOPOLOGY 為 2x2x2,則產品為 8。由於 tpu7x-standard-4t 中的每個 VM 都有 4 個晶片,因此節點數量為 2。
    • spot:將使用 Spot VM 做為 TPU 節點的節點集區。 節點集區建立後即無法變更這項設定。詳情請參閱「Spot VM」。
    • flex_start:將使用彈性啟動用量方案的節點集區。如果已啟用 spot,這項設定就無法設為 true。

其他 TPU 版本

您可以使用 Google Cloud CLI、Terraform 或 Google Cloud 控制台,在第 3、4、5p、5e 和 Trillium (第 6e 版) 中建立多主機 TPU 節點集區。

gcloud

  gcloud container node-pools create POOL_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --cluster=CLUSTER_NAME \
      --node-locations=NODE_ZONE \
      --machine-type=MACHINE_TYPE \
      --tpu-topology=TPU_TOPOLOGY \
      [--num-nodes=NUM_NODES] \
      [--spot \]
      [--flex-start \]
      [--enable-autoscaling \
        --max-nodes MAX_NODES]
      [--reservation-affinity=specific \
      --reservation=RESERVATION_NAME] \
      [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
      [--placement-type=COMPACT]

更改下列內容:

  • POOL_NAME:新節點集區的名稱。
  • CONTROL_PLANE_LOCATION:根據您要使用的 TPU 版本,選擇適用的可用區名稱。如要找出可用位置,請參閱「GKE 中的 TPU 可用性」。
  • CLUSTER_NAME:叢集名稱。
  • NODE_ZONES:以半形逗號分隔的一或多個區域清單,GKE 會在這些區域中建立節點集區。

    選用:使用 AI 區域,例如 us-central1-ai1a。AI 可用區是專門的地點,經過最佳化調整,可處理 Google Cloud 區域內的 AI/機器學習工作負載。

  • MACHINE_TYPE:用於節點的機器類型。如要進一步瞭解可用的機型,請參閱「選擇 TPU 版本」。
  • TPU_TOPOLOGY:TPU 配量的實體拓撲。拓撲格式取決於 TPU 版本。如要進一步瞭解 TPU 拓撲,請參閱「選擇拓撲」一文中的表格。

    詳情請參閱「拓撲」。

    您也可以視需要使用下列旗標:

  • NUM_NODES:節點集區中的節點數量。這個值必須是零,或是 TPU_TOPOLOGY ({A}x{B}x{C}) 中定義的值除以每個 VM 中的晶片數。如果是多主機 TPU v4 和 TPU v5e,每個 VM 中的晶片數量為四個。因此,如果 TPU_TOPOLOGY 為 2x4x4 (每個 VM 有四顆晶片的 TPU v4),則 NUM_NODES 為 32/4,等於 8。如果省略這個旗標,系統會根據拓撲和機型計算節點數量,並預設為該數量。

  • RESERVATION_NAME:GKE 在建立節點集區時使用的預留名稱。如果省略這個旗標,GKE 會使用可用的 TPU 配量節點集區。如要進一步瞭解 TPU 預留項目,請參閱「TPU 預留項目」。

  • --spot:將節點集區設為使用 Spot VM,用於 TPU 節點。節點集區建立後即無法變更這項設定。詳情請參閱「Spot VM」。

  • --flex-start:將節點集區設為使用彈性啟動型 VM。彈性啟動 VM 是使用 flex-start 用量方案建立,GKE 1.33.0-gke.1712000 以上版本支援這項方案。

  • --enable-autoscaling:建立已啟用自動調度資源的節點集區。當 GKE 擴充多主機 TPU 配量節點集區時,會以不可分割的方式將節點集區從零擴充至最大大小。

    • MAX_NODES:節點集區的大小上限。如果提供 --enable-autoscaling,則必須提供 --max-nodes 旗標,且該旗標必須等於 TPU_TOPOLOGY ({A}x{B}x{C}) 中定義的值除以每個 VM 中的晶片數量。
  • --node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY:告知 GKE 多主機 TPU 配量節點集區為集合。如果符合下列條件,請使用這個標記:

    • 節點集區會執行推論工作負載。
    • 節點集區使用 TPU Trillium。
    • Spot VM 不支援集合排程。

    如要進一步瞭解集合排程管理,請參閱「在多主機 TPU 切片中管理集合排程」。

  • --placement-type=COMPACT:建立節點集區時啟用密集配置。 這個選項必須搭配 --tpu-topology 旗標使用。 詳情請參閱「建立密集配置政策」和「TPU 拓撲」。

Terraform

  1. 請確認您使用的是 4.84.0 以上版本的 google 供應商。
  2. 在 Terraform 設定中新增下列區塊:

    resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
      provider           = google
      project            = PROJECT_ID
      cluster            = CLUSTER_NAME
      name               = POOL_NAME
      location           = CLUSTER_LOCATION
      node_locations     = [NODE_ZONES]
      initial_node_count = NUM_NODES
    
      autoscaling {
        max_node_count = MAX_NODES
        location_policy      = "ANY"
      }
      node_config {
        machine_type = MACHINE_TYPE
        reservation_affinity {
          consume_reservation_type = "SPECIFIC_RESERVATION"
          key = "compute.googleapis.com/reservation-name"
          values = [RESERVATION_LABEL_VALUES]
        }
        flex_start = false
        spot = true
      }
    
      placement_policy {
        type = "COMPACT"
        tpu_topology = TPU_TOPOLOGY
      }
    }
    

    更改下列內容:

    • NODE_POOL_RESOURCE_NAME:Terraform 範本中的節點集區資源名稱。
    • PROJECT_ID:您的專案 ID。
    • CLUSTER_NAME:要新增節點集區的現有叢集名稱。
    • POOL_NAME:要建立的節點集區名稱。
    • CLUSTER_LOCATION:叢集的運算位置。建議您使用區域性叢集,提高 Kubernetes 控制層的可靠性。您也可以使用區域叢集。 詳情請參閱「選取 TPU 版本和拓撲」。
    • NODE_ZONES:以半形逗號分隔的一或多個區域清單,GKE 會在這些區域中建立節點集區。

      選用:使用 AI 區域,例如 us-central1-ai1a。AI 可用區是專門的地點,經過最佳化調整,可處理 Google Cloud 區域內的 AI/機器學習工作負載。

    • NUM_NODES:節點集區中的節點數量。這個值必須為零,或是 TPU 晶片數量除以四的結果,因為在多主機 TPU 配量中,每個 TPU 配量節點都有 4 個晶片。舉例來說,如果 TPU_TOPOLOGY 是 4x8,則代表有 32 個晶片,因此 NUM_NODES 必須是 8。如要進一步瞭解 TPU 拓撲,請參閱「選擇 TPU 版本」一文中的表格。
    • TPU_TOPOLOGY:這表示 TPU 節點的實體拓撲。拓撲格式取決於您使用的 TPU 版本。如要進一步瞭解 TPU 拓撲,請參閱「選擇拓撲」一文中的表格。

    您也可以視需要使用下列變數:

    • RESERVATION_NAME:如果您使用 TPU 預留項目,這是建立節點集區時要使用的預留項目資源標籤清單。如要進一步瞭解如何填入「RESERVATION_LABEL_VALUES」欄位中的「reservation_affinity」,請參閱 Terraform 供應商。
    • autoscaling:建立已啟用自動調度資源的節點集區。當 GKE 擴充多主機 TPU 配量節點集區時,會不可分割地將節點集區從零擴充至最大大小。
      • MAX_NODES:節點集區的大小上限。這個值必須等於 TPU_TOPOLOGY ({A}x{B}x{C}) 中定義的值的乘積,除以每個 VM 中的晶片數量。
    • spot:允許節點集區將 Spot VM 用於 TPU 節點。 節點集區建立後即無法變更這項設定。詳情請參閱「Spot VM」。
    • flex_start:將節點集區設為使用彈性啟動用量方案。如果已啟用 spot,則無法設為 true。

控制台

如要建立含 TPU 的節點集區,請按照下列步驟操作:

  1. 前往 Google Cloud 控制台的「Google Kubernetes Engine」頁面。

    前往「Google Kubernetes Engine」

  2. 在叢集清單中,按一下您要修改的叢集名稱。

  3. 按一下「Nodes」(節點) 分頁標籤。

  4. 按一下「建立由使用者管理的節點集區」。

  5. 在「節點集區詳細資料」部分中,勾選「指定節點位置」方塊。

  6. 根據要使用的 TPU 版本選取可用區名稱。如要找出可用位置,請參閱「GKE 中的 TPU 可用性」。

  7. 按一下導覽窗格中的「Nodes」(節點)。

  8. 在「機器設定」部分,選取「TPU」。

  9. 在「系列」下拉式選單中,選取下列任一選項:

    • CT3:TPU v3,單一主機裝置
    • CT3P:TPU v3,多主機 Pod 配量
    • CT4P:TPU v4
    • CT5LP:TPU v5e
    • CT5P:TPU v5p
    • CT6E:TPU Trillium (v6e)
  10. 在「機型」下拉式選單中,選取要用於節點的機器名稱。請參閱「選擇 TPU 版本」表格,瞭解如何定義機型和 TPU 拓撲,以建立多主機 TPU 配量節點集區。

  11. 在「TPU 拓撲」下拉式選單中,選取 TPU 配量的實體拓撲。

  12. 在「需要變更」對話方塊中,按一下「進行變更」。

  13. 確認「開機磁碟類型」為「標準永久磁碟」或「SSD 永久磁碟」。

  14. 或者,您也可以勾選「啟用 Spot VM 上的節點」核取方塊,在節點集區中使用 Spot VM 做為節點。

  15. 點選「建立」。

使用 GKE 自動佈建節點功能

您可以設定 GKE 自動建立及刪除節點集區,以滿足 TPU 工作負載的資源需求。

  1. 如要啟用節點集區自動佈建功能,請編輯叢集 TPU 資源限制:

      gcloud container clusters update CLUSTER_NAME \
          --location=CONTROL_PLANE_LOCATION \
          --enable-autoprovisioning \
          --min-cpu=MINIMUM_CPU \
          --min-memory=MINIMUM_MEMORY \
          --max-cpu=MAXIMUM_CPU \
          --max-memory=MAXIMUM_MEMORY \
          --min-accelerator=type=TPU_TYPE,count=MINIMUM_TPU_COUNT \
          --max-accelerator=type=TPU_TYPE,count=MAXIMUM_TPU_COUNT
    

    更改下列內容:

    • TPU_TYPE:TPU 類型。舉例來說,使用 tpu7x-standard-4t 代表 Ironwood (TPU7x)。
    • MINIMUM_TPU_COUNT:叢集可擁有的指定類型 TPU 晶片數量下限。如果您指定的值大於多主機 TPU 配量中的 TPU 晶片數量,GKE 會移除配量中的所有節點。多主機節點集區的規模介於 0 到切片中的節點數量之間,沒有中間值。
    • MAXIMUM_TPU_COUNT:叢集可擁有的指定類型 TPU 晶片數量上限。如果是多主機 TPU 節點,請指定大於每個節點晶片數的值,讓 GKE 能以原子方式調度節點。配量中的晶片數量是 TPU 拓撲的乘積。舉例來說,如果拓撲為 2x2x2,切片中的晶片數量為 8,表示 MAXIMUM_TPU_COUNT 的值必須大於 8。

定義自訂 ComputeClass

您也可以使用自訂 ComputeClass,將 GKE 設定為在建立新節點的資源調度作業期間要求 TPU。

您可以在自訂 ComputeClass 規格中指定 TPU 設定選項。當 GKE 工作負載使用該自訂 ComputeClass 時,GKE 會嘗試在擴充時,佈建使用您指定設定的 TPU。

以下各節說明如何建立自訂 ComputeClass,然後建立使用 ComputeClass 中定義 TPU 的 Job。

建立自訂 ComputeClass

建立遵循 TPU 規則的自訂 ComputeClass 時,步驟會因您使用 Ironwood (TPU7x) 或較早的 TPU 版本而有所不同。

Ironwood (TPU7x)

  1. 建立工作負載政策。只有在建立多主機節點集區時,才需要執行這個步驟,這取決於您選擇的拓撲。如果使用單一主機節點集區,請略過這個步驟。

    gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
        --type=HIGH_THROUGHPUT \
        --accelerator-topology=TPU_TOPOLOGY \
        --project=PROJECT_ID \
        --region=REGION
    

    更改下列內容:

    • WORKLOAD_POLICY_NAME:工作負載政策的名稱。
    • TPU_TOPOLOGY:TPU Ironwood (TPU7x) 拓撲。例如,使用 2x2x2。如要進一步瞭解所有支援的 Ironwood (TPU7x) 拓撲,請參閱拓撲部分。
    • PROJECT_ID: Google Cloud 專案 ID。
    • REGION:工作負載政策的區域。工作負載政策是區域性資源,可在節點集區之間共用。
  2. 將下列資訊清單儲存為 tpu-compute-class.yaml:

    apiVersion: cloud.google.com/v1
    kind: ComputeClass
    metadata:
      name: tpu-class
    spec:
      priorities:
        - tpu:
            type: tpu7x
            topology: TPU_TOPOLOGY
            count: 4
          placement:
            policyName: WORKLOAD_POLICY_NAME
      nodePoolAutoCreation:
        enabled: true
    
  3. (選用) 您可以消耗特定預留項目或子區塊。舉例來說,您可以在 ComputeClass 資訊清單中新增下列 specs:

      reservations:
        affinity: Specific
        specific:
          - name: RESERVATION_NAME
            reservationBlock:
              name: RESERVATION_BLOCK_NAME
              reservationSubBlock:
                name: RESERVATION_SUB_BLOCK_NAME
    

    更改下列內容:

    • RESERVATION_NAME:Compute Engine 容量預留項目的名稱。
    • RESERVATION_BLOCK_NAME:Compute Engine 容量預留區塊的名稱。
    • RESERVATION_SUB_BLOCK_NAME:Compute Engine 容量預留子區塊的名稱。

    詳情請參閱「耗用預留的可用區資源」。

其他 TPU 版本

如要使用為 TPU 設定的自訂 ComputeClass,佈建 v3、v4、v5p、v5e 或 v6e (Trillium) TPU,請完成下列步驟:

  1. 將下列資訊清單儲存為 tpu-compute-class.yaml:

    apiVersion: cloud.google.com/v1
    kind: ComputeClass
    metadata:
      name: tpu-class
    spec:
      priorities:
      - tpu:
          type: TPU_TYPE
          count: NUMBER_OF_CHIPS
          topology: TOPOLOGY
      - spot: true
        tpu:
          type: TPU_TYPE
          count: NUMBER_OF_CHIPS
          topology: TOPOLOGY
      - flexStart:
          enabled: true
        tpu:
          type: TPU_TYPE
          count: NUMBER_OF_CHIPS
          topology: TOPOLOGY
      nodePoolAutoCreation:
        enabled: true
    

    更改下列內容:

    • TPU_TYPE:要使用的 TPU 類型,例如 tpu-v4-podslice。必須是 GKE 支援的值。
    • TOPOLOGY:配量中的 TPU 晶片排列方式,例如 2x2x4。必須是所選 TPU 類型支援的拓撲。
    • NUMBER_OF_CHIPS:容器使用的 TPU 晶片數量。limits 和 requests 的值必須相同。
  2. 部署 ComputeClass:

    kubectl apply -f tpu-compute-class.yaml
    

    如要進一步瞭解自訂 ComputeClass 和 TPU,請參閱「TPU 設定」。

建立使用 TPU 的 Job

  1. 將下列資訊清單儲存為 tpu-job.yaml:

    apiVersion: v1
    kind: Service
    metadata:
      name: headless-svc
    spec:
      clusterIP: None
      selector:
        job-name: tpu-job
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: tpu-job
    spec:
      backoffLimit: 0
      completions: 4
      parallelism: 4
      completionMode: Indexed
      template:
        spec:
          subdomain: headless-svc
          restartPolicy: Never
          nodeSelector:
            cloud.google.com/compute-class: tpu-class
          containers:
          - name: tpu-job
            image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
            ports:
            - containerPort: 8471 # Default port using which TPU VMs communicate
            - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
            command:
            - bash
            - -c
            - |
              python -c 'import jax; print("TPU cores:", jax.device_count())'
            resources:
              requests:
                cpu: 10
                memory: MEMORY_SIZE
                google.com/tpu: NUMBER_OF_CHIPS
              limits:
                cpu: 10
                memory: MEMORY_SIZE
                google.com/tpu: NUMBER_OF_CHIPS
    

    更改下列內容:

    • NUMBER_OF_CHIPS:容器使用的 TPU 晶片數量。limits 和 requests 的值必須相同,且等於所選自訂 ComputeClass 中的 CHIP_COUNT 值。
    • MEMORY_SIZE:TPU 使用的記憶體量上限。記憶體限制取決於您使用的 TPU 版本和拓撲。詳情請參閱「加速器的最低和最高限制」。
    • NUMBER_OF_CHIPS:容器使用的 TPU 晶片數量。limits 和 requests 的值必須相同。
  2. 部署工作:

    kubectl create -f tpu-job.yaml
    

    建立這項 Job 時,GKE 會自動執行下列操作:

    • 佈建節點來執行 Pod。視您指定的 TPU 類型、拓撲和資源要求而定,這些節點可以是單一主機配量或多主機配量。視最高優先順序的 TPU 資源可用性而定,GKE 可能會改用較低的優先順序,盡可能提高加速器容量的彈性、效率和可用性。
    • 在 Pod 中新增 taint,並在節點中新增容許條件,防止其他工作負載與 TPU 工作負載在相同節點上執行。

    詳情請參閱「關於自訂 ComputeClass」。

  3. 完成本節後,您可以刪除建立的資源,避免繼續計費:

    kubectl delete -f tpu-job.yaml
    

準備工作負載

TPU 工作負載必須符合下列準備條件。

  1. JAX、PyTorch 和 TensorFlow 等架構會使用libtpu共用程式庫存取 TPU VM。libtpu包含 XLA 編譯器、TPU 執行階段軟體和 TPU 驅動程式。每個版本的 PyTorch 和 JAX 都需要特定 libtpu.so 版本。為避免套件版本衝突,建議使用 JAX AI 圖片。如要在 GKE 中使用 TPU,請務必使用下列版本:tpu7x
    TPU 類型 「libtpu.so」版本
    Ironwood (TPU7x)
    TPU Trillium (v6e)
    tpu-v6e-slice
    TPU v5e
    tpu-v5-lite-podslice
    TPU v5p
    tpu-v5p-slice
    TPU v4
    tpu-v4-podslice
    TPU v3
    tpu-v3-slice
    tpu-v3-device
  2. 在工作負載資訊清單中新增 Kubernetes 節點選取器,確保 GKE 會在您定義的 TPU 機型和 TPU 拓撲上排定 TPU 工作負載:

      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
        cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
        cloud.google.com/placement-policy-name: WORKLOAD_POLICY # Required only for Ironwood (TPU7x)
      

    更改下列內容:

    • TPU_ACCELERATOR:TPU 加速器的名稱。例如,使用 tpu7x-standard-4t。
    • TPU_TOPOLOGY:TPU 配量的實體拓撲。拓撲格式取決於 TPU 版本。例如,使用 2x2x2。如要瞭解詳情,請參閱「規劃 GKE 中的 TPU」。
    • WORKLOAD_POLICY:您要用來放置 TPU Pod 的工作負載政策名稱。只有 Ironwood (TPU7x) 需要這個節點選取器。

完成工作負載準備後,即可執行使用 TPU 的 Job。

以下各節提供範例,說明如何執行使用 TPU 執行基本運算的 Job。

在 TPU 節點上執行工作負載

本節說明如何準備工作負載,以及如何執行工作負載的範例。

範例 1:執行 Deployment,在 Pod 規格中要求 TPU

gcloud

GKE 會使用 Pod 或 ComputeClass 中的設定,判斷 TPU 節點的設定。下列資訊清單是 Deployment 規格的範例,要求 Pod 規格中的 TPU。如果已啟用叢集層級的自動佈建節點設定,這項部署作業會觸發節點集區自動建立程序。建立這個範例 Deployment 時,GKE 會建立節點集區,其中包含具有 2x2x2 拓撲和兩部 ct4p-hightpu-4t 機器的 TPU v4 區塊。

apiVersion: apps/v1
kind: Deployment
metadata:
  name: tpu-workload
  labels:
    app: tpu-workload
spec:
  replicas: 2
  template:
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x2
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        securityContext:
          privileged: true # Required for GKE versions earlier than 1.28 to access TPUs.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("Total TPU chips:", jax.device_count())'
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4
        ports:
        - containerPort: 80

在這個資訊清單中,下列欄位定義了 TPU 設定:

  • cloud.google.com/gke-tpu-accelerator:TPU 版本和類型。舉例來說,使用 tpu7x-standard-4t 代表 Ironwood (TPU7x)。
  • cloud.google.com/gke-tpu-topology:拓撲,包含 TPU 晶片的數量和實體排列方式。例如,使用 2x2x2。
  • limits.google.com/tpu:每個 VM 的 TPU 晶片數量。舉例來說,如果您使用 tpu7x-standard-4t,則每個 VM 的 TPU 晶片數量為 4。

Cluster Toolkit

使用 gcluster job submit 指令提交工作負載。如要進一步瞭解工作提交選項,請參閱「Cluster Toolkit Job Submission Guide」:

gcluster job submit \
  --name=WORKLOAD_NAME \
  --cluster=DEPLOYMENT_NAME \
  --project=PROJECT_ID \
  --location=ZONE \
  --compute-type=COMPUTE_TYPE \
  --topology=TOPOLOGY \
  --num-slices=1 \
  --image=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest \
  --command="python -c 'import jax; print(\"TPU cores:\", jax.device_count())'"

更改下列內容:

  • WORKLOAD_NAME:工作負載的名稱。
  • DEPLOYMENT_NAME:Cluster Toolkit 部署作業的名稱 (也是 GKE 叢集的名稱)。
  • PROJECT_ID:您的 Google Cloud 專案 ID。
  • ZONE:根據您要使用的 TPU 版本,選擇適用的可用區名稱。如要找出可用位置,請參閱「GKE 中的 TPU 可用性」。
  • COMPUTE_TYPE:與節點集區加速器類型相符的運算類型,例如 v6e-16、tpu7x-standard-4t 或 v5e-8。
  • TOPOLOGY:與節點集區相符的實體拓撲,例如 2x4 或 2x2x2。

範例 2:執行工作負載,顯示 TPU 配量節點集區中可用的 TPU 晶片數量

下列工作負載會傳回多主機 TPU 切片中所有節點的 TPU 晶片數量。如要建立多主機切片,工作負載必須具備下列參數:

  • TPU 版本:TPU v4
  • 拓撲:2x2x4

選取這個版本和拓撲會產生多主機配量。

  1. 將下列資訊清單儲存為 available-chips-multihost.yaml:
    apiVersion: v1
    kind: Service
    metadata:
      name: headless-svc
    spec:
      clusterIP: None
      selector:
        job-name: tpu-available-chips
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: tpu-available-chips
    spec:
      backoffLimit: 0
      completions: 4
      parallelism: 4
      completionMode: Indexed
      template:
        spec:
          subdomain: headless-svc
          restartPolicy: Never
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice # Node selector to target TPU v4 slice nodes.
            cloud.google.com/gke-tpu-topology: 2x2x4 # Specifies the physical topology for the TPU slice.
          containers:
          - name: tpu-job
            image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
            ports:
            - containerPort: 8471 # Default port using which TPU VMs communicate
            - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
            securityContext:
              privileged: true # Required for GKE versions earlier than 1.28 to access TPUs.
            command:
            - bash
            - -c
            - |
              python -c 'import jax; print("TPU cores:", jax.device_count())' # Python command to count available TPU chips.
            resources:
              requests:
                cpu: 10
                memory: 407Gi
                google.com/tpu: 4 # Request 4 TPU chips for this workload.
              limits:
                cpu: 10
                memory: 407Gi
                google.com/tpu: 4 # Limit to 4 TPU chips for this workload.
  2. 部署資訊清單:
    kubectl create -f available-chips-multihost.yaml
    

    GKE 會執行含有四個 VM 的 TPU v4 配量 (多主機 TPU 配量)。配量有 16 個互連的 TPU 晶片。

  3. 確認 Job 已建立四個 Pod:
    kubectl get pods
    

    輸出結果會與下列內容相似:

    NAME                       READY   STATUS      RESTARTS   AGE
    tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
    tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
    tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
    tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
    
  4. 取得其中一個 Pod 的記錄檔:
    kubectl logs POD_NAME
    

    將 POD_NAME 替換為其中一個已建立的 Pod 名稱。例如:tpu-job-podslice-0-5cd8r。

    輸出結果大致如下:

    TPU cores: 16
    
  5. 選用:移除工作負載:
    kubectl delete -f available-chips-multihost.yaml
    

範例 3:執行工作負載,顯示 TPU 配量中可用的 TPU 晶片數量

下列工作負載是靜態 Pod,會顯示附加至特定節點的 TPU 晶片數量。如要建立單一主機節點,工作負載必須具備下列參數:

  • TPU 版本:TPU v5e
  • 拓撲:2x4

這個版本和拓撲選取項目會產生單一主機切片。

  1. 將下列資訊清單儲存為 available-chips-singlehost.yaml:
    apiVersion: v1
    kind: Pod
    metadata:
      name: tpu-job-jax-v5
    spec:
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice # Node selector to target TPU v5e slice nodes.
        cloud.google.com/gke-tpu-topology: 2x4 # Specify the physical topology for the TPU slice.
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        securityContext:
          privileged: true # Required for GKE versions earlier than 1.28 to access TPUs.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("Total TPU chips:", jax.device_count())'
        resources:
          requests:
            google.com/tpu: 8 # Request 8 TPU chips for this container.
          limits:
            google.com/tpu: 8 # Limit to 8 TPU chips for this container.
  2. 部署資訊清單:
    kubectl create -f available-chips-singlehost.yaml
    

    GKE 會佈建節點,其中包含八個使用 TPU v5e 的單一主機 TPU 配量。每個 TPU 節點有八個 TPU 晶片 (單主機 TPU 配量)。

  3. 取得 Pod 的記錄:
    kubectl logs tpu-job-jax-v5
    

    輸出結果大致如下:

    Total TPU chips: 8
    
  4. 選用:移除工作負載:
      kubectl delete -f available-chips-singlehost.yaml
      

範例 4:執行以兩個 NUMA 節點為目標的工作負載

為提升工作負載效能,GKE 可讓您在多容器設定中部署工作負載。這會將每個容器繫結至特定 NUMA 節點內的 CPU、記憶體和 TPU 資源。如要進一步瞭解 NUMA 繫結,請參閱「GKE 中的 Ironwood (TPU7x) 簡介」。

  1. 建立名為 topology-manager-config.yaml 的檔案,並在其中加入下列內容:

    kubeletConfig:
      cpuManagerPolicy: static
      topologyManager:
        policy: best-effort
        scope: container
    

    這項設定會將容器要求的資源對齊單一 NUMA 節點。GKE 必須進行這項設定,才能接納多容器 TPU 工作負載,並正確指派晶片。如需更多政策選項,請參閱「Topology Manager Policies」。

  2. 如要將 kubelet 設定套用至節點集區,請建立新的節點集區或更新現有節點集區。這個節點集區必須使用 Ironwood (TPU7x) 節點,例如 tpu7x-standard-4t 機型,才能與下一個步驟中的節點選取器相符。

    • 如要建立啟用拓撲管理工具的新節點集區,請按照下列步驟操作:

      gcloud container node-pools create NODE_POOL_NAME \
          --location=LOCATION \
          --cluster=CLUSTER_NAME \
          --node-locations=NODE_ZONES \
          --machine-type=tpu7x-standard-4t \
          --system-config-from-file=topology-manager-config.yaml
      

      更改下列內容:

      • NODE_POOL_NAME:新節點集區的名稱。
      • LOCATION:叢集區域。
      • CLUSTER_NAME:叢集名稱。
      • NODE_ZONES:以半形逗號分隔的清單,列出一或多個節點集區的可用區。
    • 如要更新已啟用拓撲管理員的現有節點集區,請按照下列步驟操作:

      gcloud container node-pools update NODE_POOL_NAME \
          --location=LOCATION \
          --cluster=CLUSTER_NAME \
          --system-config-from-file=topology-manager-config.yaml
      

      更改下列內容:

      • NODE_POOL_NAME:節點集區名稱。
      • LOCATION:叢集區域。
      • CLUSTER_NAME:叢集名稱。
  3. 建立名為 tpu-jobset.yaml 的檔案,並加入下列資訊清單:

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: tpu7x-test
    spec:
      failurePolicy:
        maxRestarts: 0
      replicatedJobs:
        - name: tpu7x
          replicas: 1
          template:
            spec:
              parallelism: 2
              completions: 2
              backoffLimit: 0
              template:
                spec:
                  subdomain: test
                  restartPolicy: Never
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu7x
                    cloud.google.com/gke-tpu-topology: 2x2x2
                  containers:
                  - name: jax-tpu-1
                    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    securityContext:
                      privileged: false
                    command:
                    - bash
                    - -c
                    - |
                      set -ex
                      python -c 'import jax; print("local devices:", jax.local_device_count(), "devices:", jax.device_count())'
                    resources:
                      requests:
                        google.com/tpu: 2
                      limits:
                        google.com/tpu: 2
                  - name: jax-tpu-2
                    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    securityContext:
                      privileged: false
                    command:
                    - bash
                    - -c
                    - |
                      set -ex
    
                      python -c 'import jax; print("local devices:", jax.local_device_count(), "devices:", jax.device_count())'
                    resources:
                      requests:
                        google.com/tpu: 2
                      limits:
                        google.com/tpu: 2
    

    這個資訊清單定義名為 tpu7x-test 的 JobSet,其中會執行兩個主要容器 jax-tpu-1 和 jax-tpu-2。每個容器會要求兩個 google.com/tpu 晶片。這兩顆晶片來自同一個 NUMA 節點。接著,每個容器都會執行 Python 指令碼,驗證 JAX 是否能存取 TPU。

    這項 JobSet 設定顯示,同一 Pod 中的多個容器可以存取及使用個別的 TPU 晶片。

  4. 套用資訊清單:

    kubectl apply -f tpu-jobset.yaml
    

範例 5:使用多 NIC Cloud Storage FUSE 磁碟區執行多容器 TPU 工作負載

這個範例擴充了範例 4,納入使用 Cloud Storage FUSE 將 Cloud Storage FUSE bucket 掛接為磁碟區的程序。這項設定會運用多 NIC 功能,進而提升網路效能。

  1. 建立名為 tpu7x-gcsfuse-multinic.yaml 的檔案,並加入下列資訊清單:

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: tpu7x-gcsfuse-multinic
    spec:
      failurePolicy:
        maxRestarts: 0
      replicatedJobs:
        - name: tpu7x
          replicas: 1
          template:
            spec:
              parallelism: 2
              completions: 2
              backoffLimit: 0
              template:
                spec:
                  subdomain: test
                  restartPolicy: Never
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu7x
                    cloud.google.com/gke-tpu-topology: 2x2x2
                  hostNetwork: true
                  volumes:
                  - name: gcs0
                    csi:
                      driver: gcsfuse.csi.storage.gke.io
                      volumeAttributes:
                        bucketName: BUCKET_NAME
                        multiNICIndex: "0"
                        mountOptions: implicit-dirs
                        hostNetworkPodKSA: "true"
                  - name: gcs1
                    csi:
                      driver: gcsfuse.csi.storage.gke.io
                      volumeAttributes:
                        bucketName: BUCKET_NAME
                        multiNICIndex: "1"
                        mountOptions: implicit-dirs
                        hostNetworkPodKSA: "true"
                  containers:
                  - name: jax-tpu-1
                    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    securityContext:
                      privileged: false
                    command:
                    - bash
                    - -c
                    - |
                      set -ex
                      python -c 'import jax; print("local devices:", jax.local_device_count(), "devices:", jax.device_count())'
                      echo "Use /gcs${WORKLOAD_NIC_PREFERRED_NUMA}/ for Cloud Storage access"
                    resources:
                      requests:
                        google.com/tpu: 2
                      limits:
                        google.com/tpu: 2
                    volumeMounts:
                    - name: gcs0
                      mountPath: /gcs0
                      readOnly: false
                    - name: gcs1
                      mountPath: /gcs1
                      readOnly: false
                  - name: jax-tpu-2
                    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    securityContext:
                      privileged: false
                    command:
                    - bash
                    - -c
                    - |
                      set -ex
    
                      python -c 'import jax; print("local devices:", jax.local_device_count(), "devices:", jax.device_count())'
                      echo "Use /gcs${WORKLOAD_NIC_PREFERRED_NUMA}/ for Cloud Storage access"
                    resources:
                      requests:
                        google.com/tpu: 2
                      limits:
                        google.com/tpu: 2
                    volumeMounts:
                    - name: gcs0
                      mountPath: /gcs0
                      readOnly: false
                    - name: gcs1
                      mountPath: /gcs1
                      readOnly: false
    

    這個資訊清單會將下列 Cloud Storage FUSE 磁碟區掛接點新增至「Example 4」。

    • multiNICIndex 會在磁碟區屬性中指定,以便為每個磁碟區使用對應 NUMA 節點上的網路介面。
    • 工作負載會使用 ${WORKLOAD_NIC_PREFERRED_NUMA} 環境變數選取 Cloud Storage FUSE 磁碟區,該磁碟區與容器指派的 NUMA 節點相關聯。GKE 會將這個變數注入 Cloud TPU 工作負載的容器。

    將 BUCKET_NAME 替換為 Cloud Storage bucket 名稱。

  2. 套用資訊清單:

    kubectl apply -f tpu7x-gcsfuse-multinic.yaml
    

使用加速器 (GPU 和 TPU) 升級節點集區

GKE 會自動升級標準叢集,包括節點集區。如要讓節點更快升級至較新版本,也可以手動升級節點集區。如要控管叢集的升級方式,請使用發布版本、維護期間和排除時段,以及推出順序。

您也可以為節點集區設定節點升級策略,例如大量升級、藍綠升級或短期升級。設定這些策略後,即可確保節點集區升級時,速度和環境中斷之間能達到最佳平衡。如果是多主機 TPU 配量節點集區,GKE 會在單一步驟中,以不可分割的方式重建整個節點集區,而不是使用設定的節點升級策略。詳情請參閱「GKE 中的 TPU 相關術語」一文中的「不可分割性」定義。

使用節點升級策略時,GKE 必須暫時佈建額外資源 (視設定而定)。如果 Google Cloud節點集區的資源容量有限 (例如嘗試使用 GPU 或 TPU 建立更多節點時,看到資源可用性錯誤),請參閱在資源受限的環境中升級。

清除所用資源

如要避免系統向您的 Google Cloud 帳戶收取本指南所用資源的費用,請考慮刪除不再排定工作負載的 TPU 節點集區。如果執行的工作負載必須正常終止,請使用 kubectl drain 清理工作負載,再刪除節點。

gcloud

刪除 TPU 節點集區:

  gcloud container node-pools delete POOL_NAME \
      --location=LOCATION \
      --cluster=CLUSTER_NAME

更改下列內容:

  • POOL_NAME:節點集區的名稱。
  • LOCATION:叢集的運算位置。
  • CLUSTER_NAME:叢集名稱。

Cluster Toolkit

刪除 Cluster Toolkit 部署作業:

  gcluster destroy DEPLOYMENT_NAME

將 DEPLOYMENT_NAME 替換為部署藍圖時建立的部署目錄名稱。

進行其他設定

以下各節說明可套用至 TPU 工作負載的其他設定。

使用多配量

您可以在 Multislice 中匯總較小的切片,處理較大的訓練工作負載。詳情請參閱「GKE 中的多切片 TPU」。

遷移 TPU 預訂

如果您有現有的 TPU 預留項目,請務必先將 TPU 預留項目遷移至新的 Compute Engine 預留項目系統。您也可以建立以 Compute Engine 為基礎的預訂系統,不需要遷移。如要瞭解如何遷移 TPU 預留項目,請參閱 TPU 預留項目。

啟用記錄功能

GKE 節點 (包括 TPU VM) 上執行的容器所產生的記錄檔,會由 GKE 記錄代理程式收集,然後傳送至 Cloud Logging,並顯示在 Cloud Logging 中。

設定 TPU 節點的自動修復功能

如果多主機 TPU 配量節點集區中的 TPU 配量節點健康狀態不良,系統會重新建立整個節點集區。而在單一主機 TPU 配量節點集區中,只有狀況不良的 TPU 節點會自動修復。

導致 TPU 配量節點健康狀態不良的條件包括:

  • 任何具有一般節點條件的 TPU 配量節點。
  • 任何 TPU 配量節點,其無法分配的 TPU 數量大於零。
  • TPU 節點中任何已停止 (因搶占) 或終止的 VM 執行個體。
  • 節點維護:如果多主機 TPU 配量節點集區中的任何 TPU 配量節點因主機維護而停機,GKE 會重新建立整個 TPU 配量節點集區。

您可以在作業記錄中查看維修狀態 (包括故障原因)。如果配額不足導致失敗,請與Google Cloud 帳戶代表聯絡,要求提高相應配額。

為 TPU 配量節點設定正常終止

在控制層執行 1.29.1-gke.1425000 以上版本的 GKE 叢集中,TPU 節點支援 SIGTERM 信號,可提醒節點即將關機。在 TPU 節點中,即將關機的通知最多可設定為 五分鐘。

如要設定 GKE 在這段通知時間內正常終止工作負載,請按照「管理 GPU 和 TPU 的 GKE 節點中斷」一文中的步驟操作。

防止排程死結

執行多個 TPU 工作或多個工作副本時,不同工作的 Pod 可以排定在同一個 TPU 節點集區中。這可能會導致排程死結,因為沒有任何作業能確保有足夠的節點來完成拓撲。

如要降低排程死結的風險,請套用下列其中一種設定:

  • 使用 Pod 反相依性規則,確保不同 Job 的 Pod 不會排程到相同的 TPU 節點集區或拓撲網域。
  • 使用 jobset.sigs.k8s.io/exclusive-topology 註解,讓 JobSet 專為自己的 Pod 宣告指定的拓撲網域 (例如節點集區)。

如需詳細的設定範例和疑難排解步驟,請參閱「排程死結」。

在沒有特殊權限模式的情況下執行容器

在 GKE 1.28 以上版本的節點中執行的容器,不需要啟用特殊權限模式即可存取 TPU。GKE 1.28 以下版本的節點需要使用權限模式。

如果 TPU 節點執行的是 1.28 以下版本,請參閱下一節:

在 TPU 配量中的 VM 上執行的容器需要存取鎖定記憶體的較高限制,因此驅動程式可以透過直接記憶體存取 (DMA) 與 TPU 晶片通訊。如要啟用這項功能,必須設定較高的 ulimit。如要縮減容器的權限範圍,請完成下列步驟:

  1. 編輯 securityContext,加入下列欄位:

    securityContext:
      capabilities:
        add: ["SYS_RESOURCE"]
    
  2. 在容器內執行下列指令,增加 ulimit,然後設定工作負載以使用 TPU 資源:

    ulimit -l 68719476736
    

如要使用 TPU v5e,叢集必須為 1.27.4-gke.900 以上版本,才能在沒有特殊權限模式的情況下執行容器。

可觀測性和指標

資訊主頁

Google Cloud 控制台中的節點集區觀測能力現已正式發布。如要查看 GKE 中 TPU 多主機節點集區的狀態,請前往 Cloud Monitoring 提供的 GKE TPU 節點集區狀態資訊主頁:

前往 GKE TPU 節點集區狀態

這個資訊主頁可提供多主機 TPU 節點集區的健康狀態洞察資料。 詳情請參閱「監控 TPU 節點和節點集區的健康狀態指標」。

在控制台的「Kubernetes Clusters」(Kubernetes 叢集) 頁面中Google Cloud ,「Observability」(可觀測性) 分頁的「Accelerators」(加速器) >「TPU」標題下方,也會顯示 TPU 可觀測性指標,例如 TPU 用量。詳情請參閱「查看可觀測性指標」。

只有在 GKE 叢集中啟用系統指標,TPU 資訊主頁才會填入資料。

執行階段指標

在 GKE 1.27.4-gke.900 以上版本中,如果 TPU 工作負載同時使用 JAX 0.4.14 以上版本,並指定 containerPort: 8431,系統就會將 TPU 使用率指標匯出為 GKE 系統指標。您可以在 Cloud Monitoring 中使用下列指標,監控 TPU 工作負載的執行階段效能:

  • 工作週期:在過去取樣期間 (60 秒) 內,TensorCore 在 TPU 晶片上積極處理作業的時間百分比。百分比越高,代表 TPU 使用率越佳。
  • 已用記憶體:分配的加速器記憶體量 (以位元組為單位)。每 60 秒取樣一次。
  • 記憶體總量:加速器記憶體總量 (以位元組為單位)。每 60 秒取樣一次。

這些指標位於 Kubernetes 節點 (k8s_node) 和 Kubernetes 容器 (k8s_container) 結構定義中。

Kubernetes 容器:

  • kubernetes.io/container/accelerator/duty_cycle
  • kubernetes.io/container/accelerator/memory_used
  • kubernetes.io/container/accelerator/memory_total

Kubernetes 節點:

  • kubernetes.io/node/accelerator/duty_cycle
  • kubernetes.io/node/accelerator/memory_used
  • kubernetes.io/node/accelerator/memory_total

監控 TPU 節點和節點集區的健康狀態指標

如果訓練工作發生錯誤或終止失敗,您可以檢查基礎架構相關指標,判斷中斷是否是由基礎節點或節點集區的問題所導致。

節點狀態

在 GKE 1.32.1-gke.1357001 以上版本中,下列 GKE 系統指標會公開 GKE 節點的狀態:

  • kubernetes.io/node/status_condition

condition 欄位會回報節點的狀況,例如 Ready、DiskPressure 和 MemoryPressure。status 欄位會顯示回報的狀況狀態,可以是 True、False 或 Unknown。這是指具有 k8s_node 受監控資源類型的指標。

這項 PromQL 查詢會顯示特定節點是否為 Ready:

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    node_name="NODE_NAME",
    condition="Ready",
    status="True"}

如要協助排解叢集問題,您可能需要查看出現其他情況的節點:

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition!="Ready",
    status="True"}

您可能想特別查看未標示 Ready 的節點:

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition="Ready",
    status="False"}

如果沒有資料,節點就已準備就緒。系統每 60 秒會取樣一次狀態條件。

您可以使用下列查詢,瞭解整個機群的節點狀態:

avg by (condition,status)(
  avg_over_time(
    kubernetes_io:node_status_condition{monitored_resource="k8s_node"}[${__interval}]))

節點集區狀態

下列GKE 系統指標適用於k8s_node_pool受監控的資源,可顯示 GKE 節點集區的狀態:

  • kubernetes.io/node_pool/status

這項指標只會針對多主機 TPU 節點集區回報。

status 欄位會回報節點集區的狀態,例如 Provisioning、Running、Error、Reconciling 或 Stopping。狀態會在 GKE API 作業完成後更新。

如要確認特定節點集區是否處於 Running 狀態,請使用下列 PromQL 查詢:

kubernetes_io:node_pool_status{
    monitored_resource="k8s_node_pool",
    cluster_name="CLUSTER_NAME",
    node_pool_name="NODE_POOL_NAME",
    status="Running"}

如要監控專案中節點集區的數量 (依狀態分組),請使用下列 PromQL 查詢:

count by (status)(
  count_over_time(
    kubernetes_io:node_pool_status{monitored_resource="k8s_node_pool"}[${__interval}]))

節點集區可用性

下列 GKE 系統指標會顯示多主機 TPU 節點集區是否可用:

  • kubernetes.io/node_pool/multi_host/available

如果節點集區中的所有節點均可使用,指標值為 True,否則為 False。指標的取樣頻率為每 60 秒一次。

如要檢查專案中多主機 TPU 節點集區的可用性,請使用下列 PromQL 查詢:

avg by (node_pool_name)(
  avg_over_time(
    kubernetes_io:node_pool_multi_host_available{
      monitored_resource="k8s_node_pool",
      cluster_name="CLUSTER_NAME"}[${__interval}]))

節點中斷次數

下列 GKE 系統指標會回報自上次取樣以來,GKE 節點的中斷次數 (每 60 秒取樣一次):

  • kubernetes.io/node/interruption_count

interruption_type (例如 TerminationEvent、MaintenanceEvent 或 PreemptionEvent) 和 interruption_reason (例如 HostError、Eviction 或 AutoRepair) 欄位可協助提供節點中斷的原因。

如要查看專案叢集中 TPU 節點的中斷情形和原因,請使用下列 PromQL 查詢:

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node"}[${__interval}]))

如要只查看主機維護事件,請更新查詢,篩選 interruption_reason 的 HW/SW Maintenance 值。使用下列 PromQL 查詢:

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node", interruption_reason="HW/SW Maintenance"}[${__interval}]))

如要查看依節點集區匯總的中斷次數,請使用下列 PromQL 查詢:

  sum by (node_pool_name,interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_pool_interruption_count{monitored_resource="k8s_node_pool", interruption_reason="HW/SW Maintenance", node_pool_name=NODE_POOL_NAME }[${__interval}]))

節點集區復原時間 (TTR)

下列 GKE 系統指標會回報 GKE 多主機 TPU 節點集區的復原期長度分布:

  • kubernetes.io/node_pool/accelerator/times_to_recover

這項指標中記錄的每個樣本,都代表節點集區從停機期間復原的單一事件。

這項指標有助於追蹤多主機 TPU 節點集區的復原時間,以及服務中斷間隔時間。

您可以使用下列 PromQL 查詢,計算叢集過去 7 天的平均修復時間 (MTTR):

sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_sum{
    monitored_resource="k8s_node_pool", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_count{
    monitored_resource="k8s_node_pool",cluster_name="CLUSTER_NAME"}[7d]))

節點集區中斷情形間隔時間 (TBI)

下列 GKE 系統指標會回報 GKE 多主機 TPU 節點集區的分布時間,也就是最近一次中斷結束到目前中斷開始之間的時間:

  • kubernetes.io/node_pool/accelerator/times_between_interruptions

這項指標記錄的每個樣本,都代表最近一次和目前中斷之間的單一時間長度。

使用下列 PromQL 查詢,計算叢集過去 7 天的中斷間隔平均時間 (MTBI):

sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_between_interruptions_sum{
    monitored_resource="k8s_node_pool", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_between_interruptions_count{
    monitored_resource="k8s_node_pool",cluster_name="CLUSTER_NAME"}[7d]))

主機指標

在 GKE 1.28.1-gke.1066000 以上版本中,TPU 節點中的 VM 會將 TPU 使用率指標匯出為 GKE 系統指標。您可以在 Cloud Monitoring 中使用下列指標,監控 TPU 主機的效能:

  • TensorCore 使用率:目前 TensorCore 的使用量百分比。TensorCore 值等於矩陣乘法單元 (MXU) 加上向量單元的總和。 TensorCore 使用率值是將過去取樣期間 (60 秒) 執行的 TensorCore 作業數量,除以相同期間支援的 TensorCore 作業數量。值越大表示使用率越高。
  • 記憶體頻寬使用率:目前所用加速器記憶體頻寬的百分比。計算方式為取樣期間 (60 秒) 內使用的記憶體頻寬,除以相同取樣期間內支援的最大頻寬。

這些指標位於 Kubernetes 節點 (k8s_node) 和 Kubernetes 容器 (k8s_container) 結構定義中。

Kubernetes 容器:

  • kubernetes.io/container/accelerator/tensorcore_utilization
  • kubernetes.io/container/accelerator/memory_bandwidth_utilization

Kubernetes 節點:

  • kubernetes.io/node/accelerator/tensorcore_utilization
  • kubernetes.io/node/accelerator/memory_bandwidth_utilization

詳情請參閱「Kubernetes 指標」和「GKE 系統指標」。

管理收藏內容排程

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

TPU Trillium 支援單一主機和多主機節點集區的集合排程,這些節點集區會執行推論工作負載。以下說明收集排程行為如何取決於您使用的 TPU 配量類型:

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

如要管理集合,請根據使用的節點集區類型執行下列任一動作。

管理多主機 TPU 配量節點集區中的收集作業排程

請使用下列工作管理多主機 TPU 配量節點集區。

  • 如要檢查多主機 TPU 節點集區是否屬於集合,請執行下列指令:

    gcloud container node-pools describe NODE_POOL_NAME \
        --location LOCATION \
        --cluster CLUSTER_NAME \
        --format="json" | jq -r \
        '"nodepool-group-name: \(.config.labels["cloud.google.com/gke-nodepool-group-name"] // "")\ngke-workload-type: \(.config.labels["cloud.google.com/gke-workload-type"] // "")"'
    

    輸出結果會與下列內容相似:

    nodepool-group-name: <code><var>NODE_POOL_COLLECTION_NAME</var></code>
    gke-workload-type: HIGH_AVAILABILITY
    

    如果多主機 TPU 配量集區是集合的一部分,輸出內容會包含下列標籤:

    • cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
    • cloud.google.com/gke-nodepool-group-name: <code><var>COLLECTION_NAME</var></code>
  • 如要取得叢集中的集合清單,請執行下列指令:

    #!/bin/bash
    
    # Replace with your cluster name, project, and location
    CLUSTER_NAME=CLUSTER_NAME
    PROJECT=PROJECT_ID
    LOCATION=LOCATION
    
    declare -A collection_names
    
    node_pools=$(gcloud container node-pools list --cluster "$CLUSTER_NAME" --project "$PROJECT" --location "$LOCATION" --format="value(name)")
    
    # Iterate over each node pool
    for pool in $node_pools; do
        # Describe the node pool and extract labels using jq
        collection_name=$(gcloud container node-pools describe "$pool" \
            --cluster "$CLUSTER_NAME" \
            --project "$PROJECT" \
            --location "$LOCATION" \
            --format="json" | jq -r '.config.labels["cloud.google.com/gke-nodepool-group-name"]')
    
        # Add the collection name to the associative array if it's not empty
        if [[ -n "$collection_name" ]]; then
            collection_names["$collection_name"]=1
        fi
    done
    
    # Print the unique node pool collection names
    echo "Unique cloud.google.com/gke-nodepool-group-name values:"
    for name in "${!collection_names[@]}"; do
        echo "$name"
    done
    

    輸出結果會與下列內容相似:

    Unique cloud.google.com/gke-nodepool-group-name values: {COLLECTION_NAME_1}, {COLLECTION_NAME_2}, {COLLECTION_NAME_3}
    
  • 如要取得集合所屬的節點集區清單,請執行下列指令:

    #!/bin/bash
    
    TARGET_COLLECTION_NAME=COLLECTION_NAME
    CLUSTER_NAME=CLUSTER_NAME
    PROJECT=PROJECT_ID
    LOCATION=LOCATION
    
    matching_node_pools=()
    
    # Get the list of all node pools in the cluster
    node_pools=$(gcloud container node-pools list --cluster "$CLUSTER_NAME" --project "$PROJECT" --location "$LOCATION" --format="value(name)")
    
    # Iterate over each node pool
    for pool in $node_pools; do
        # Get the value of the cloud.google.com/gke-nodepool-group-name label
        collection_name=$(gcloud container node-pools describe "$pool" \
            --cluster "$CLUSTER_NAME" \
            --project "$PROJECT" \
            --location "$LOCATION" \
            --format="json" | jq -r '.config.labels["cloud.google.com/gke-nodepool-group-name"]')
    
        # Check if the group name matches the target value
        if [[ "$collection_name" == "$TARGET_COLLECTION_NAME" ]]; then
            matching_node_pools+=("$pool")
        fi
    done
    
    # Print the list of matching node pools
    echo "Node pools with collection name '$TARGET_COLLECTION_NAME':"
    for pool in "${matching_node_pools[@]}"; do
        echo "$pool"
    done
    

    輸出結果會與下列內容相似:

    Node pools with collection name 'COLLECTION_NAME':
    {NODE_POOL_NAME_1}
    {NODE_POOL_NAME_2}
    {NODE_POOL_NAME_3}
    
  • 如要擴大集合,請建立另一個多主機 TPU 配量節點集區,並新增 cloud.google.com/gke-workload-type 和 cloud.google.com/gke-nodepool-group-name。在 cloud.google.com/gke-nodepool-group-name 中使用相同的集合名稱,並執行相同的工作負載類型。如果叢集已啟用自動佈建節點功能,GKE 會根據工作負載需求自動建立集區。

  • 如要縮減集合,請刪除節點集區。

  • 如要刪除集合,請移除所有附加的節點集區。您可以刪除節點集區或刪除叢集。刪除叢集會一併移除其中的所有集合。

管理單主機 TPU 配量節點集區中的收集作業排程

使用下列工作管理單一主機 TPU 節點集區。

  • 如要檢查單一主機 TPU 節點集區是否已啟用集合排程,請執行下列指令:

    gcloud container node-pools describe NODE_POOL_NAME \
        --cluster CLUSTER_NAME \
        --project PROJECT_NAME \
        --location LOCATION \
        --format="json" | jq -r '.config.labels["cloud.google.com/gke-workload-type"]'
    

    輸出結果會與下列內容相似:

    gke-workload-type: HIGH_AVAILABILITY
    

    如果單一主機 TPU 配量集區屬於集合,輸出內容會包含 cloud.google.com/gke-workload-type: HIGH_AVAILABILITY 標籤。

  • 如要擴大集合,請手動或自動調整節點集區大小,並使用節點自動佈建功能。

  • 如要縮減集合,請刪除節點集區。

  • 如要刪除集合,請移除所有附加的節點集區。您可以刪除節點集區或刪除叢集。刪除叢集會一併移除其中的所有集合。

已知問題

  • 在這些節點回報可用的 TPU 之前,叢集自動調整程式可能會錯誤計算新 TPU 節點的容量。叢集自動配置器可能會執行額外的擴充作業,導致建立的節點數量超出需求。叢集自動配置器會在一般縮減作業後,縮減不需要的額外節點。
  • 如果 TPU 節點集區的節點切片處於等待狀態超過 10 小時,叢集自動調度器就會取消擴充作業。叢集自動配置器稍後會重試這類擴充作業。這可能會導致未使用預留項目的客戶無法使用 TPU。
  • 如果非 TPU 工作負載容許 TPU 汙點,且在排空 TPU 配量節點集區期間重新建立,可能會導致節點集區無法縮減。
  • 記憶體頻寬使用率指標不適用於 v5e TPU。

後續步驟