使用 TPU 和彈性啟動 VM 執行小型批次工作負載

本指南說明如何使用彈性啟動 VM,為中小型訓練工作負載最佳化 TPU 佈建作業。彈性啟動 VM 是使用 flex-start 用量方案建立。在本指南中,您會使用彈性啟動 VM 部署工作負載,其中包含 TPU 節點集區。

本指南適用於機器學習 (ML) 工程師、平台管理員和操作員,以及有興趣使用 Kubernetes 容器自動化調度管理功能執行批次工作負載的資料和 AI 專家。如要進一步瞭解 Google Cloud 內容中提及的常見角色和範例工作,請參閱「常見的 GKE 使用者角色和工作」。

彈性啟動定價

如果工作負載需要視需要動態佈建資源,且短期預訂最多七天,不需要複雜的配額管理,並希望以具成本效益的方式存取資源,建議使用彈性啟動。彈性啟動功能採用 Dynamic Workload Scheduler,並按照Dynamic Workload Scheduler 定價計費:

  • vCPU、GPU 和 TPU 最高可享 47% 的折扣。
  • 您只需支付實際用量費用。

事前準備

開始之前,請確認您已完成下列工作:

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

建立採用彈性啟動的節點集區

如果您使用 Autopilot 模式的叢集,請略過本節,並前往「執行批次工作負載」一節。

如要在現有標準叢集上建立啟用彈性啟動功能的節點集區,請使用 gcloud CLI。

建立單一主機 TPU 節點集區

您可以使用彈性啟動功能建立單一主機 TPU 配量節點集區:

  1. 建立採用彈性啟動模式的節點集區:

    gcloud container node-pools create NODE_POOL_NAME \
        --cluster=CLUSTER_NAME \
        --location=CONTROL_PLANE_LOCATION \
        --node-locations=NODE_ZONES \
        --machine-type=MACHINE_TYPE \
        --reservation-affinity=none \
        --enable-autoscaling \
        --flex-start \
        --num-nodes 0 \
        --min-nodes=0 \
        --max-nodes=1
    

    更改下列內容:

    • NODE_POOL_NAME:您為節點集區選擇的名稱。
    • CLUSTER_NAME:叢集名稱。
    • CONTROL_PLANE_LOCATION:叢集控制層的運算區域
    • NODE_ZONES:以半形逗號分隔的清單,列出 GKE 建立節點集區的一或多個區域。
    • MACHINE_TYPE:節點使用的機器類型。如要進一步瞭解與 TPU 相容的機器類型,請參閱「選擇 TPU 版本」一文中的表格。

      舉例來說,節點集區建立指令可以包含下列參數:

      ...
      --machine-type=ct6e-standard-4t \
      --tpu-topology=4x4 \
      --enable-autoscaling \
      --num-nodes=0 \
      --max-nodes=4 \
      

      這個指令會將 --max-nodes 欄位設為 4,因為 4x4 拓撲包含 16 個晶片,而每個 ct6e-standard-4t VM 有 4 個晶片。

      叢集自動配置器會擴充至工作負載所需的節點數量。工作負載完成後,叢集自動配置器會將節點數縮減為零。

    • --reservation-affinity=none彈性啟動不會使用或需要預留項目。

建立多主機 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 \
        --location=us-central1 \
        --node-locations=us-central1-c \
        --machine-type=tpu7x-standard-4t \
        --reservation-affinity=none \
        --enable-autoscaling \
        --num-nodes=0 --min-nodes=0 --max-nodes=MAX_NODES \
        --flex-start \
        --placement-policy=WORKLOAD_POLICY
    

    更改下列內容:

    • NODE_POOL_NAME:新節點集區的名稱。
    • WORKLOAD_POLICY:您建立的 Workload 政策名稱。
    • MAX_NODES:節點集區的大小上限。如果提供 --enable-autoscaling,則必須提供 --max-nodes 旗標,且該旗標必須等於 TPU_TOPOLOGY ({A}x{B}x{C}) 中定義的值除以每個 VM 中的晶片數。舉例來說,如果 TPU_TOPOLOGY2x2x2,則產品為 8。由於 tpu7x-standard-4t 中的每個 VM 都有 4 個晶片,因此節點數為 2。

    這個指令會建立名為 NODE_POOL_NAME 的節點集區,並具備下列特徵:

    • --machine-type=tpu7x-standard-4t 會指定 Ironwood (TPU7x) 機型。
    • --flex-start 可啟用彈性啟動。

Terraform

  1. 請確認您使用的是 google 4.84.0 以上版本。
  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 = true
      }
    
      placement_policy {
        policy_name = WORKLOAD_POLICY_NAME
      }
    }
    

    更改下列內容:

    • NODE_POOL_RESOURCE_NAME:Terraform 範本中的節點集區資源名稱。
    • PROJECT_ID:您的專案 ID。
    • CLUSTER_NAME:要新增節點集區的現有叢集名稱。
    • POOL_NAME:要建立的節點集區名稱。
    • NODE_ZONES:以半形逗號分隔的清單,列出 GKE 建立節點集區的一或多個區域。
    • NUM_NODES:節點集區中的節點數量。這個值必須是零,或是 TPU 晶片數量除以四的結果,因為在多主機 TPU 配量中,每個 TPU 配量節點都有四個晶片。舉例來說,如果 TPU_TOPOLOGY4x8,則代表有 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 中的晶片數。舉例來說,如果 TPU_TOPOLOGY2x2x2,則產品為 8。由於 tpu7x-standard-4t 中的每個 VM 都有 4 個晶片,因此節點數為 2。
    • spot:將使用 Spot VM 做為 TPU 配量節點的節點集區。 節點集區建立後,這項設定就無法再變更。詳情請參閱「Spot VM」。
    • flex_start:將使用彈性啟動消耗選項的節點集區。如果已啟用 spot,這項設定就無法設為 true

其他 TPU 版本

您可以使用 Google Cloud CLI、Terraform 或 Google Cloud 控制台,在 v3、v4、v5p、v5e 和 Trillium (v6e) 版本中建立多主機 TPU 節點集區。

gcloud

  gcloud container node-pools create NODE_POOL_NAME \
      --cluster=CLUSTER_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --node-locations=NODE_ZONES \
      --machine-type=MACHINE_TYPE \
      --tpu-topology=TPU_TOPOLOGY \
      --reservation-affinity=none \
      --enable-autoscaling \
      --num-nodes=0 --min-nodes=0 --max-nodes=MAX_NODES \
      --flex-start

更改下列內容:

  • NODE_POOL_NAME:新節點集區的名稱。
  • CLUSTER_NAME:叢集名稱。
  • CONTROL_PLANE_LOCATION:根據您要使用的 TPU 版本,指定區域名稱。如要找出可用位置,請參閱「GKE 中的 TPU 可用性」。
  • NODE_ZONES:以半形逗號分隔的清單,列出 GKE 建立節點集區的一或多個區域。
  • MACHINE_TYPE:節點使用的機器類型。如要進一步瞭解與 TPU 相容的機器類型,請參閱「選擇 TPU 版本」一文中的表格。
  • TPU_TOPOLOGY:TPU 拓撲。 例如,2x2x2。如要查看所有支援的 TPU 拓撲,請參閱拓撲部分
  • MAX_NODES:節點集區的大小上限。如果提供 --enable-autoscaling,則必須提供 --max-nodes 旗標,且該旗標必須等於 TPU_TOPOLOGY ({A}x{B}x{C}) 中定義的值除以每個 VM 中的晶片數。

    這個指令會建立名為 NODE_POOL_NAME 的節點集區,並啟用彈性啟動功能。

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]
        }
        flex_start = 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 建立節點集區的一或多個區域。
    • NUM_NODES:節點集區中的節點數量。這個值必須是零,或是 TPU 晶片數量除以四的結果,因為在多主機 TPU 配量中,每個 TPU 配量節點都有 4 個晶片。舉例來說,如果 TPU_TOPOLOGY4x8,則代表有 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. 在「系列」下拉式選單中,選取下列其中一個選項:

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

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

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

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

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

  14. 點選「建立」

驗證節點集區中的彈性啟動狀態

執行下列指令:

gcloud container node-pools describe NODE_POOL_NAME \
    --cluster CLUSTER_NAME \
    --location CONTROL_PLANE_LOCATION \
    --format="get(config.flexStart)"

如果節點集區已啟用彈性啟動功能,則 flexStart 欄位會設為 True

執行批次工作負載

在本節中,您將建立 Job,排定使用彈性啟動 VM 的 TPU 節點。 Kubernetes 中的 Job 控制器會建立一或多個 Pod,並確保這些 Pod 成功執行特定工作。

  1. Google Cloud 控制台中,按一下「啟用 Cloud Shell」,啟動 Cloud Shell 工作階段。Cloud Shell 啟用圖示工作階段會在 Google Cloud 控制台的底部窗格中開啟。

  2. 建立名為 dws-flex-start.yaml 的檔案:

    單一主機

    請使用下列資訊清單來建立 dws-flex-start.yaml 檔案:

    apiVersion: batch/v1
    kind: Job
    metadata:
      name: job-1
    spec:
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-flex-start: "true"
            cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
            cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
          containers:
          - name: container-1
            image: gcr.io/k8s-staging-perf-tests/sleep:latest
            args: ["3600s"] # Sleep for 1 hour
            resources:
              requests:
                  google.com/tpu: NUM_CHIPS
              limits:
                  google.com/tpu: NUM_CHIPS
          restartPolicy: OnFailure
    

    多主機

    請使用下列資訊清單來建立 dws-flex-start.yaml 檔案:

    apiVersion: v1
    kind: Service
    metadata:
      name: headless-svc
    spec:
      clusterIP: None
      selector:
        job-name: job-1
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: job-1
    spec:
      backoffLimit: 0
      completions: 2
      parallelism: 2
      completionMode: Indexed
      template:
        spec:
          subdomain: headless-svc
          restartPolicy: Never
          nodeSelector:
              cloud.google.com/gke-flex-start: "true"
              cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
              cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
          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
            command:
            - bash
            - -c
            - |
              python -c 'import jax; print("TPU cores:", jax.device_count())'
            resources:
              requests:
                google.com/tpu: NUM_CHIPS
              limits:
                google.com/tpu: NUM_CHIPS
    

    更改下列內容:

    • ACCELERATOR_TYPE:建立節點集區時使用的 TPU 加速器類型。例如 tpu-v4-podslicetpu-v5-lite-podslice
    • TPU_TOPOLOGY:TPU 節點的實體拓撲。舉例來說,視 TPU 版本而定,值可能是 4x4x42x2
    • NUM_CHIPS:每個 VM 中的 TPU 晶片數量為一、四或八。詳情請參閱「TPU 版本」。
  3. 套用 dws-flex-start.yaml 資訊清單:

    kubectl apply -f dws-flex-start.yaml
    
  4. 確認 Job 在同一節點上執行:

    kubectl get pods
    

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

    NAME    READY   STATUS      RESTARTS   AGE   IP       NODE               NOMINATED NODE   READINESS GATES
    job-1   0/1     Completed   0          19m   10.(...) gke-flex-zonal-a2  <none>           <none>
    

清除所用資源

如要避免系統向您的 Google Cloud 帳戶收取本頁面所用資源的費用,請刪除含有相關資源的專案,或者保留專案但刪除個別資源。

刪除專案

  1. In the Google Cloud console, go to the Manage resources page.

    Go to Manage resources

  2. In the project list, select the project that you want to delete, and then click Delete.
  3. In the dialog, type the project ID, and then click Shut down to delete the project.

刪除個別資源

  1. 刪除工作:

    kubectl delete job -l "job-name in (job-1,job-2)"
    
  2. 刪除節點集區:

    gcloud container node-pools delete NODE_POOL_NAME \
          --location CONTROL_PLANE_LOCATION
    
  3. 刪除叢集:

    gcloud container clusters delete CLUSTER_NAME
    

後續步驟