使用 TPU 和彈性啟動佈建模式執行小批次工作負載


本指南說明如何使用彈性啟動佈建模式,針對中小型訓練工作負載,最佳化 TPU 佈建作業。在本指南中,您會使用彈性啟動模式,部署由 TPU 節點集區組成的節點。

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

彈性啟動定價

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

  • vCPU、GPU 和 TPU 可享折扣 (最高 53%)。
  • 即付即用。

事前準備

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

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

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

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

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

您可以透過彈性啟動建立單一或多主機 TPU 配量節點集區:

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

    單一主機

    gcloud container node-pools create NODE_POOL_NAME \
        --cluster=CLUSTER_NAME \
        --location=LOCATION_NAME \
        --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:叢集名稱。
    • LOCATION_NAME:叢集控制層的運算區域
    • NODE_ZONES:以半形逗號分隔的清單,列出 GKE 建立節點集區的一或多個區域。
    • MACHINE_TYPE:節點使用的機器類型。如要進一步瞭解與 TPU 相容的機器類型,請參閱「選擇 TPU 版本」一文中的表格。

    多主機

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

    更改下列內容:

    • NODE_POOL_NAME:您為節點集區選擇的名稱。
    • CLUSTER_NAME:叢集名稱。
    • LOCATION_NAME:叢集控制層的運算區域
    • NODE_ZONES:以半形逗號分隔的清單,列出 GKE 建立節點集區的一或多個區域。
    • MACHINE_TYPE:節點使用的機器類型。舉例來說,您可以將 ct6e-standard-4t 用於 TPU Trillium。如要進一步瞭解可用的機器類型,請參閱「選擇 TPU 版本」。
    • TPU_TOPOLOGY:TPU 配量的實體拓撲。拓撲格式取決於 TPU 版本。如要進一步瞭解 TPU 拓撲,請參閱「選擇拓撲」一文中的表格。

    使用彈性啟動建立節點集區時,上述指令會使用下列必要旗標:

    • --enable-autoscalingflex-start 只會在工作負載執行時,提供必要的運算資源。您必須設定下列參數:

      • --num-nodes=0
      • --min-nodes=0
      • --max-nodes 設為 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彈性啟動不會使用或要求預留項目。

  2. 確認節點集區的彈性啟動狀態:

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

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

執行批次工作負載

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

  1. Trusted Cloud 控制台中,按一下「Activate Cloud Shell」(啟用 Cloud Shell),啟動 Cloud Shell 工作階段。Cloud Shell 啟用圖示工作階段會在 Trusted 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: python:3.10
            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
            - |
              pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
              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>
    

清除所用資源

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

刪除專案

  1. In the Trusted 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 LOCATION_NAME
    
  3. 刪除叢集:

    gcloud container clusters delete CLUSTER_NAME
    

後續步驟