在 GKE 中部署 TPU 多切片

本頁說明如何使用 Cloud TPU Multislice 設定，在 Google Kubernetes Engine (GKE) 中部署工作負載，以經濟實惠的方式進行大規模訓練。

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

在 GKE 中設定 Multislice 之前，請務必熟悉下列概念：

  1. Cloud TPU 簡介
  2. Cloud TPU 系統架構
  3. 關於 GKE 中的 TPU

什麼是 TPU 多配量？

TPU 多配量是 TPU 配量中 VM 的架構組織，其中兩個以上的 Cloud TPU 配量會透過資料中心網路 (DCN) 通訊。Multislice 可提供完整堆疊，以符合成本效益的方式大規模訓練，近線性擴充至數萬個 TPU 晶片。在多配量設定中，GKE 會在多個 TPU 配量上部署多配量工作負載。配量內的 TPU 晶片會透過晶片間互連 (ICI) 進行通訊。切片之間的通訊會透過 DCN 進行。

如果 Job 過大，無法放在單一 TPU 節點上，建議使用 Multislice。

GKE 中的多切片可用性

  • 標準版支援 1.27.4-gke.900 以上版本的 Multislice。
  • Autopilot 支援 1.29.2-gke.1521000 以上版本的 Multislice。
  • Multislice 支援 JAX 和 PyTorch 架構。支援的最低 JAX 版本為 2.1。
  • 多配量僅支援多主機 TPU 配量節點集區。舉例來說，您無法搭配 ct4p-hightpu-4t 拓撲使用 Multislice，因為這些是單一主機 TPU 節點集區。2x2x1ct5lp-hightpu-4t2x2
  • Multislice 僅支援同步多控制器訓練。
  • Multislice 工作負載只能在共用相同 TPU 類型、大小和拓撲的 TPU Slice 上執行。
  • 多配量不支援 TPU v3。

事前準備

開始之前，請確認你已完成下列工作：

  • 啟用 Google Kubernetes Engine API。
    • 啟用 Google Kubernetes Engine API
  • 如要使用 Google Cloud CLI 執行這項工作，請安裝並初始化 gcloud CLI。如果您先前已安裝 gcloud CLI，請執行 gcloud components update，取得最新版本。

在 Multislice 上執行工作負載

本節說明如何在 Multislice 上執行工作負載。如果您使用 GKE Autopilot 模式，請跳至「執行 Multislice 工作負載」一節。執行 1.29.2-gke.1521000 以上版本的 Autopilot 叢集預設會啟用 TPU。

準備標準模式節點集區

本節將說明下列步驟：

  1. 建立三個多主機 TPU 配量節點集區
  2. 驗證節點集區狀態

建立 TPU 節點集區

您可以建立多個多主機 TPU 配量節點集區。在本指南中，請建立三個多主機 TPU 配量節點集區，以執行多配量工作負載。您可以使用 Google Cloud CLI、Terraform 或 Google Cloud 控制台，建立多主機 TPU 節點集區。

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --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：新節點集區的名稱。
  • LOCATION：可用區名稱，取決於您要使用的 TPU 版本。如要找出可用位置，請參閱「GKE 中的 TPU 可用性」。
  • CLUSTER_NAME：叢集名稱。
  • NODE_ZONES：以半形逗號分隔的清單，列出 GKE 建立節點集區的一或多個可用區。
  • 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：將節點集區設為使用彈性啟動佈建模式。

  • --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. 請確認您使用的是 google 4.84.0 以上版本。

  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]
    }
    spot = true
    flex_start = false
  }

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

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

    • RESERVATION_NAME：如果您使用 TPU 預留，這是建立節點集區時要使用的預留資源標籤清單。如要進一步瞭解如何填入 reservation_affinity 欄位中的 RESERVATION_LABEL_VALUES，請參閱 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. 按一下 「Add node pool」(新增節點集區)。

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

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

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

  7. 在「機器設定」專區中，選取「TPU」。

  8. 在「系列」下拉式選單中，選取下列其中一個選項：

    • CT3P：適用於 TPU v3。
    • CT4P：適用於 TPU v4。
    • CT5LP：適用於 TPU v5e。

  9. 在「Machine type」(機器類型) 下拉式選單中，選取要用於節點的機器名稱。請參閱「選擇 TPU 版本」表格，瞭解如何定義機器類型和 TPU 拓撲，以建立多主機 TPU 配量節點集區。

  10. 在「TPU Topology」(TPU 拓撲) 下拉式選單中，選取 TPU 配量的實體拓撲。

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

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

  13. 如要為節點集區中的節點使用 Spot VM，請勾選「啟用 Spot VM 上的節點」核取方塊。

  14. 點選「建立」。

驗證節點集區狀態

  1. 取得憑證，以便使用 kubectl 存取叢集：

    gcloud container clusters get-credentials CLUSTER_NAME \
    --project=PROJECT_ID \
    --location=CONTROL_PLANE_LOCATION

    更改下列內容：

    • CLUSTER_NAME：叢集名稱。
    • PROJECT_ID：您的專案 ID。
    • CONTROL_PLANE_LOCATION：叢集控制層的 Compute Engine 位置。為地區叢集提供地區，或為區域叢集提供區域。

  2. 在 Cloud Shell 中使用 kubectl，查看 TPU 配量節點：

    kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=ACCELERATOR_TYPE \
   -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY

    更改下列內容：

    • TPU_ACCELERATOR：建立節點集區時使用的 TPU 加速器類型。例如 tpu-v4-podslice 或 tpu-v5-lite-podslice。
    • TPU_TOPOLOGY：TPU 配量的實體拓撲。

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

     NAME                                    STATUS   ROLES    AGE    VERSION
 gke-tpu-20ee2cce-5tv6                   Ready    <none>   34h     v1.28.1-gke.1066000

執行 Multislice 工作負載

在本節中，您會執行 JAX 工作負載，顯示 TPU 配量中的全域 TPU 晶片數量，然後結束。

如要執行 JAX 工作負載，請按照下列步驟操作：

  1. 建立下列 tpu-multislice.yaml 資訊清單：

    Autopilot

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: NUM_SLICES
      template:
        spec:
          parallelism: NUM_NODES
          completions: NUM_NODES
          backoffLimit: 0
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
              containers:
              - name: jax-tpu
                image: python:3.10
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                 limits:
                    google.com/tpu: NUM_CHIPS

    標準

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: NUM_SLICES
      template:
        spec:
          parallelism: NUM_NODES
          completions: NUM_NODES
          backoffLimit: 0
          template:
            spec:
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
              containers:
              - name: jax-tpu
                image: python:3.10
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                  limits:
                   google.com/tpu: NUM_CHIPS

    更改下列內容：

    • NUM_SLICES：TPU 配量節點集區的數量。在本例中，NUM_SLICES 等於 3。
    • ACCELERATOR_TYPE：建立節點集區時使用的 TPU 加速器類型。例如 tpu-v4-podslice 或 tpu-v5-lite-podslice。
    • TPU_TOPOLOGY：TPU 配量的實體拓撲。例如 4x4x4 或 2x2 (視 TPU 版本而定)。
    • NUM_NODES：節點集區中的節點數量。 這個值必須為零，或是 TPU_TOPOLOGY ({A}x{B}x{C}) 中定義的值的乘積，除以每個 VM 中的 TPU 晶片數量。如果是多主機 TPU v4，每個 VM 的 TPU 晶片數量為四個。如果是多主機 TPU v5e，每個 VM 中的 TPU 晶片數量為一、四或八個。因此，如果 TPU_TOPOLOGY 為 2x4x4 (每個 VM 中有四個 TPU 晶片的 TPU v4)，則 NUM_NODES 為 32/4，等於 8。
    • NUM_CHIPS：如果是多主機 TPU v4，每個 VM 中的 TPU 晶片數量為四個。如果是多主機 TPU v5e，每個 VM 中的 TPU 晶片數量為一、四或八個。詳情請參閱「TPU 配量中 VM 上的 TPU 晶片」。

    在這個資訊清單中：

    • JobSet 是無頭服務，名稱與 JobSet 名稱相同，在本例中為 multislice-job。
    • alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool 註解會設定 Pod 相依性，確保所有 Pod「安排至同一區塊」。
    • maxRestarts: 4 表示子項 Job 失敗時，GKE 重新啟動 JobSet 的次數上限。如果 JobSet 重新啟動次數達到定義上限，JobSet 就會標示為失敗。
    • parallelism 和 completions 欄位等於每個節點集區中的節點數量。
    • backoff 為 0，因為 Multislice 僅支援同步多控制器訓練。必須設為 0。只要有任何 Pod 失敗，工作就會失敗。
    • 親和性區段中的值可確保一組多重切片中只執行一個 TPU 多重切片工作負載。
    • containerPort: 8080 是 MXLA 協調器的連接埠
    • containerPort: 8431 是匯出 TPU 使用率指標的連接埠
    • securityContext: privileged: true 表示節點已啟用具備存取 TPU 權限的模式。在 GKE 1.28 以上版本中，節點不需要啟用特殊權限模式，即可存取 TPU。詳情請參閱「Run containers without privileged mode」。

  2. 套用資訊清單：

    kubectl apply -f tpu-multislice.yaml

  3. 確認工作負載已獲准：

    kubectl get jobsets

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

    NAME            RESTARTS   COMPLETED   AGE
multislice-job                         3s

  4. 監控已佈建 Pod 的狀態：

    kubectl get pods

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

     NAME                                READY   STATUS      RESTARTS   AGE
 multislice-job-slice-0-0-wzq9t      0/1     Completed   0          2m31s
 multislice-job-slice-0-1-zf4dp      0/1     Completed   0          2m30s
 multislice-job-slice-1-0-hbfn5      0/1     Completed   0          2m31s
 multislice-job-slice-1-1-45fgl      0/1     Completed   0          2m30s
 multislice-job-slice-2-0-wjbp4      0/1     Completed   0          2m30s
 multislice-job-slice-2-1-lwnvs      0/1     Completed   0          2m30s

    multislice-job JobSet 會排定、建立及執行 Pod，直到完成為止。Pod 名稱的格式為 <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>。jobsetName 前置字元會決定 Pod 所屬的 JobSet。

  5. 選用：移除 JAX 工作負載：

    kubectl delete -f tpu-multislice.yaml

進行其他設定

以下各節說明可套用至 Multislice 的其他設定。

使用 hostNetwork 提升網路效能

為提升 TPU 節點間的網路效能，建議您開啟 hostNetworking。在 Pod 規格中使用 hostNetwork: true，即可略過所有 Kubernetes 網路堆疊，讓 Kubernetes Pod 直接使用主機網路進行 VM 對 VM 通訊。

如要開啟 hostNetworking，請在 Pod 規格中新增以下兩行：

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

如要繼續使用 podHostnames 搭配 hostNetwork 探索工作站節點，請設定 dnsPolicy: ClusterFirstWithHostNet。如果您執行自動繼續訓練作業，且需要使用相同名稱重新載入相同檢查點，這項功能就非常重要。

如果您使用 TPU Trillium (v6e)，且 Pod 使用 hostNetworking，請安裝下列 DaemonSet，在節點上調整 /proc/sys/net/ipv4/tcp_rmem。

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/51bf3dcab6ff658cf62cc32867f96860bf58dfdc/scripts/network-setup/v6e-increase-rmem.yaml

在 TPU Trillium 上改善網路效能，而不使用 hostNetwork

如果您使用 TPU Trillium，但 Pod 無法使用 hostNetworking，請啟用 netdevice 模式的多重網路功能，以獲得最佳網路效能。netdevice 模式 NIC 支援多網路，可將 VM NIC 直接傳遞至 Pod，略過 Kubernetes 和 GKE Dataplane V2。

ct6e-standard-4t 機型由兩個實體 NIC 支援。Kubernetes 需要一個無法傳遞至 Pod 的 vNIC。因此，每個節點都必須有三個 vNIC，才能讓 Pod 直接存取兩個 vNIC，進而發揮兩個實體 NIC 的最佳效能。

如要為 ct6e-standard-4t 啟用 netdevice 模式，請完成下列步驟：

  1. 建立另外兩個支援 netdevice 模式的虛擬私有雲
  2. 建立具備多網路功能的 GKE 叢集

  3. 設定兩個 netdevice 網路。舉例來說，您可以使用下列 GKENetworkParamSet 和 Network 物件 (SECOND_VPC 和 THIRD_VPC 是在上一個步驟中建立的虛擬私有雲)：

    apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: tpu-second
spec:
  vpc: SECOND_VPC
  vpcSubnet: SECOND_VPC_SUBNET
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: tpu-third
spec:
  vpc: THIRD_VPC
  vpcSubnet: SECOND_VPC_SUBNET
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: tpu-second
spec:
  provider: "GKE"
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: tpu-second
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: tpu-third
spec:
  provider: "GKE"
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: tpu-third

  4. 使用三個網路連線 Pod。舉例來說，您可以在 Pod 規格中使用下列註解：

    metadata:
  annotations:
    networking.gke.io/default-interface: 'eth0'
    networking.gke.io/interfaces: |
      [
        {"interfaceName":"eth0","network":"default"},
        {"interfaceName":"eth1","network":"tpu-second"},
        {"interfaceName":"eth2","network":"tpu-third"},
      ]

  5. 在 Pod 內、在 init 容器或應用程式容器中套用網路 sysctl。舉例來說，您可以在 Pod 規格中新增下列 init 容器：

    initContainers:
- name: "network-optimization-sysctls"
  image: "busybox"
  securityContext:
    privileged: true
  command:
  - bash
  - -c
  - |
    echo 5000 > /proc/sys/net/ipv4/tcp_rto_min_us
    echo 1 > /proc/sys/net/ipv4/tcp_no_metrics_save
    echo 0 > /proc/sys/net/ipv4/tcp_slow_start_after_idle
    echo 131072 > /proc/sys/net/core/optmem_max
    echo "4096 41943040 314572800" > /proc/sys/net/ipv4/tcp_rmem
最佳做法：

請使用 eth1 和 eth2 介面提升網路效能，而非 eth0 介面。只要在工作負載規格中加入 export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" 即可。

啟用記錄功能

如果叢集已啟用 GKE 系統記錄，您就能在記錄檔探索工具中，查看 GKE 節點 (包括 TPU 配量節點) 上執行的容器所發出的記錄。

您可以使用記錄檔探索工具，透過下列篩選條件查看工作負載的容器記錄，從 GKE 查看記錄：

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME

請使用下列篩選器，查看 TPU 節點和工作站：

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>

詳情請參閱「查看 GKE TPU 記錄檔」。

啟用其他指標

除了常見的 TPU 指標，還有 4 個多切片專屬的 TPU 執行階段指標。這些指標適用於 GKE 1.29.1-gke.1016000 以上版本。TPU 工作負載必須使用 JAX 0.4.24 版0.4.24

可用的多切片指標如下：

  • DCN (資料中心網路) 傳輸延遲：多切片流量的網路傳輸延遲分配情形。
  • 集體延遲：多重切片流量的端對端集體延遲分佈情形。
  • 主機到裝置的傳輸延遲：多切片流量中，每個資料區塊的主機到裝置傳輸延遲分配情形。
  • 裝置到主機的傳輸延遲：多切片流量中，每個資料區塊的裝置到主機傳輸延遲分配情形。

這些指標位於 Kubernetes 容器 (k8s_container) 結構定義中：

  • kubernetes.io/container/multislice/network/dcn_transfer_latencies
  • kubernetes.io/container/multislice/network/collective_end_to_end_latencies
  • kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies
  • kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies

TPU 配量與多配量

下表說明 TPU 節點和多節點的架構組織差異：

TPU 配量 多配量
互連性 工作負載會在單一 TPU 切片上執行。配量中的所有 TPU 晶片都會透過 ICI 連線。 工作負載會在多個 TPU 配量上執行。切片內的通訊是透過 ICI 進行。切片之間的通訊是透過 DCN 進行。
支援的節點集區 單一主機 TPU 配量和多主機 TPU 配量 多主機 TPU 配量群組
建議的工作負載類型 IndexedJob 或 JobSet JobSet

清除所用資源

如要避免付費，最簡單的方法就是刪除您為了本教學課程所建立的專案。 Google Cloud 或者，您也可以刪除個別資源。

刪除專案

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

刪除個別資源

刪除 GKE 叢集：

```sh
gcloud container clusters delete  CLUSTER_NAME \
   --project=PROJECT_ID  \
   --location=CONTROL_PLANE_LOCATION
```

後續步驟