針對混合式 AI/機器學習訓練和推論工作負載，盡可能提高 GKE 資源使用率

本教學課程說明如何在單一 Google Kubernetes Engine (GKE) 叢集中，有效率地在訓練和推論服務工作負載之間共用加速器資源。將混合工作負載分配到單一叢集，可提高資源使用率、簡化叢集管理作業、減少因加速器數量限制而發生的問題，並提升整體成本效益。

在本教學課程中，您將使用 Gemma 2 大型語言模型 (LLM) 進行推論，並使用 Hugging Face TGI (文字生成介面) 服務架構，建立高優先順序的服務部署作業，同時建立低優先順序的 LLM 微調作業。這兩項工作負載都會在單一叢集上執行，且該叢集使用 NVIDIA L4 GPU。您可以使用 Kueue (Kubernetes 原生的開放原始碼工作佇列系統)，管理及排定工作負載。Kueue 可讓您優先處理服務工作，並先佔用低優先順序訓練工作，以最佳化資源使用率。隨著服務需求減少，您重新分配釋出的加速器，以繼續訓練作業。您可以使用 Kueue 和優先順序類別，在整個過程中管理資源配額。

本教學課程適用於機器學習 (ML) 工程師、平台管理員和作業人員，以及想在 GKE 叢集上訓練及代管機器學習 (ML) 模型，並降低成本和管理負擔的資料和 AI 專家，尤其是在處理數量有限的加速器時。如要進一步瞭解我們在 Google Cloud 內容中提及的常見角色和範例工作，請參閱「常見的 GKE 使用者角色和工作」。

閱讀本頁面之前，請先熟悉下列概念：

準備環境

在本節中，您將佈建部署 TGI 和模型所需的資源，以用於推論和訓練工作負載。

取得模型存取權

如要存取 Gemma 模型，並部署至 GKE，請先簽署授權同意聲明協議，然後產生 Hugging Face 存取權杖。

  1. 簽署授權同意聲明協議。前往模型同意聲明頁面，使用 Hugging Face 帳戶驗證同意聲明，然後接受模型條款。

  2. 產生存取權杖。如要透過 Hugging Face 存取模型，您需要 Hugging Face 權杖。如要產生新權杖 (如果沒有)，請按照下列步驟操作：

    1. 依序點選「Your Profile」(你的個人資料) >「Settings」(設定) >「Access Tokens」(存取權杖)。
    2. 選取「New Token」。
    3. 指定您選擇的名稱，以及至少 Read 的角色。
    4. 選取「產生權杖」。
    5. 將產生的權杖複製到剪貼簿。

啟動 Cloud Shell

在本教學課程中，您將使用 Cloud Shell 管理Google Cloud上託管的資源。Cloud Shell 已預先安裝本教學課程所需的軟體，包括 kubectl、 gcloud CLI 和 Terraform。

如要使用 Cloud Shell 設定環境，請按照下列步驟操作：

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

  2. 設定預設環境變數：

    gcloud config set project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)

    將 PROJECT_ID 替換為您的 Google Cloud 專案 ID。

  3. 從 GitHub 複製範例程式碼。在 Cloud Shell 中執行下列指令：

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/
cd kubernetes-engine-samples/ai-ml/mix-train-and-inference
export EXAMPLE_HOME=$(pwd)

建立 GKE 叢集

您可以將 Autopilot 或標準叢集用於混合工作負載。建議您使用 Autopilot 叢集，享受全代管的 Kubernetes 體驗。如要選擇最適合工作負載的 GKE 作業模式，請參閱「選擇 GKE 作業模式」。

Autopilot

  1. 在 Cloud Shell 中設定預設環境變數：

    export HF_TOKEN=HF_TOKEN
export REGION=REGION
export CLUSTER_NAME="llm-cluster"
export PROJECT_NUMBER=$(gcloud projects list \
    --filter="$(gcloud config get-value project)" \
    --format="value(PROJECT_NUMBER)")
export MODEL_BUCKET="model-bucket-$PROJECT_ID"

    替換下列值：

    • HF_TOKEN：您先前產生的 Hugging Face 權杖。
    • REGION：支援您要使用的加速器類型，例如 L4 GPU 的 us-central1。

    您可以調整 MODEL_BUCKET 變數，這個變數代表您儲存訓練模型權重的 Cloud Storage 值區。

  2. 建立 Autopilot 叢集：

    gcloud container clusters create-auto ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${REGION} \
    --release-channel=rapid

  3. 為微調作業建立 Cloud Storage bucket：

    gcloud storage buckets create gs://${MODEL_BUCKET} \
    --location ${REGION} \
    --uniform-bucket-level-access

  4. 如要授予 Cloud Storage 值區的存取權，請執行下列指令：

    gcloud storage buckets add-iam-policy-binding "gs://$MODEL_BUCKET" \
    --role=roles/storage.objectAdmin \
    --member=principal://iam.googleapis.com/projects/$PROJECT_NUMBER/locations/global/workloadIdentityPools/$PROJECT_ID.svc.id.goog/subject/ns/llm/sa/default \
    --condition=None

  5. 如要取得叢集的驗證憑證，請執行下列指令：

    gcloud container clusters get-credentials llm-cluster \
    --location=$REGION \
    --project=$PROJECT_ID

  6. 為 Deployment 建立命名空間。在 Cloud Shell 中執行下列指令：

    kubectl create ns llm

標準

  1. 在 Cloud Shell 中設定預設環境變數：

    export HF_TOKEN=HF_TOKEN
export REGION=REGION
export CLUSTER_NAME="llm-cluster"
export GPU_POOL_MACHINE_TYPE="g2-standard-24"
export GPU_POOL_ACCELERATOR_TYPE="nvidia-l4"
export PROJECT_NUMBER=$(gcloud projects list \
    --filter="$(gcloud config get-value project)" \
    --format="value(PROJECT_NUMBER)")
export MODEL_BUCKET="model-bucket-$PROJECT_ID"

    替換下列值：

    • HF_TOKEN：您先前產生的 Hugging Face 權杖。
    • REGION：支援您要使用的加速器類型，例如 L4 GPU 的 us-central1。

    您可以調整下列變數：

    • GPU_POOL_MACHINE_TYPE：要在所選區域中使用的節點集區機器系列。這個值取決於您選取的加速器類型。詳情請參閱「在 GKE 上使用 GPU 的限制」。舉例來說，本教學課程使用 g2-standard-24，每個節點都連接兩個 GPU。如需最新可用 GPU 清單，請參閱「適用於運算工作負載的 GPU」。
    • GPU_POOL_ACCELERATOR_TYPE：所選區域支援的加速器類型。舉例來說，本教學課程使用 nvidia-l4。如需最新可用 GPU 清單，請參閱「適用於運算工作負載的 GPU」。
    • MODEL_BUCKET：儲存訓練模型權重的 Cloud Storage 值區。

  2. 建立標準叢集：

    gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${REGION} \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --release-channel=rapid \
    --machine-type=e2-standard-4 \
    --addons GcsFuseCsiDriver \
    --num-nodes=1

  3. 為推論和微調工作負載建立 GPU 節點集區：

    gcloud container node-pools create gpupool \
    --accelerator type=${GPU_POOL_ACCELERATOR_TYPE},count=2,gpu-driver-version=latest \
    --project=${PROJECT_ID} \
    --location=${REGION} \
    --node-locations=${REGION}-a \
    --cluster=${CLUSTER_NAME} \
    --machine-type=${GPU_POOL_MACHINE_TYPE} \
    --num-nodes=3

  4. 為微調作業建立 Cloud Storage bucket：

    gcloud storage buckets create gs://${MODEL_BUCKET} \
    --location ${REGION} \
    --uniform-bucket-level-access

  5. 如要授予 Cloud Storage 值區的存取權，請執行下列指令：

    gcloud storage buckets add-iam-policy-binding "gs://$MODEL_BUCKET" \
    --role=roles/storage.objectAdmin \
    --member=principal://iam.googleapis.com/projects/$PROJECT_NUMBER/locations/global/workloadIdentityPools/$PROJECT_ID.svc.id.goog/subject/ns/llm/sa/default \
    --condition=None

  6. 如要取得叢集的驗證憑證，請執行下列指令：

    gcloud container clusters get-credentials llm-cluster \
    --location=$REGION \
    --project=$PROJECT_ID

  7. 為 Deployment 建立命名空間。在 Cloud Shell 中執行下列指令：

    kubectl create ns llm

為 Hugging Face 憑證建立 Kubernetes 密鑰

如要建立包含 Hugging Face 權杖的 Kubernetes Secret，請執行下列指令：

kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=$HF_TOKEN \
    --dry-run=client -o yaml | kubectl apply --namespace=llm --filename=-

設定 Kueue

在本教學課程中，Kueue 是中央資源管理工具，可讓訓練和服務工作負載有效共用 GPU。Kueue 會定義資源需求 (「規格」)、透過佇列排定工作負載優先順序 (以服務工作優先於訓練工作)，並根據需求和優先順序動態分配資源，藉此達成上述目標。本教學課程使用 Workload 資源類型，分別將推論和微調工作負載分組。

Kueue 的搶占功能可確保高優先順序的服務工作負載一律有必要資源，方法是在資源不足時暫停或逐出低優先順序的訓練工作。

如要使用 Kueue 控制推論伺服器 Deployment，請啟用 pod 整合，並設定 managedJobsNamespaceSelector 排除 kube-system 和 kueue-system 命名空間。

  1. 在 /kueue 目錄中，查看 kustomization.yaml 中的程式碼。這個資訊清單會安裝 Kueue 資源管理工具，並採用自訂設定。

    apiVersion: kustomize.config.k8s.io/v1beta1
kind: Kustomization
resources:
- https://github.com/kubernetes-sigs/kueue/releases/download/v0.12.3/manifests.yaml
patches:
- path: patch.yaml
  target:
    version: v1
    kind: ConfigMap
    name: kueue-manager-config

  2. 在 /kueue 目錄中，查看 patch.yaml 中的程式碼。這個 ConfigMap 會自訂 Kueue，排除管理 kube-system 和 kueue-system 命名空間中的 Pod。

    apiVersion: v1
kind: ConfigMap
metadata:
  name: kueue-manager-config
data:
  controller_manager_config.yaml: |
    apiVersion: config.kueue.x-k8s.io/v1beta1
    kind: Configuration
    health:
      healthProbeBindAddress: :8081
    metrics:
      bindAddress: :8080
    # enableClusterQueueResources: true
    webhook:
      port: 9443
    leaderElection:
      leaderElect: true
      resourceName: c1f6bfd2.kueue.x-k8s.io
    controller:
      groupKindConcurrency:
        Job.batch: 5
        Pod: 5
        Workload.kueue.x-k8s.io: 5
        LocalQueue.kueue.x-k8s.io: 1
        ClusterQueue.kueue.x-k8s.io: 1
        ResourceFlavor.kueue.x-k8s.io: 1
    clientConnection:
      qps: 50
      burst: 100
    #pprofBindAddress: :8083
    #waitForPodsReady:
    #  enable: false
    #  timeout: 5m
    #  blockAdmission: false
    #  requeuingStrategy:
    #    timestamp: Eviction
    #    backoffLimitCount: null # null indicates infinite requeuing
    #    backoffBaseSeconds: 60
    #    backoffMaxSeconds: 3600
    #manageJobsWithoutQueueName: true
    managedJobsNamespaceSelector:
      matchExpressions:
        - key: kubernetes.io/metadata.name
          operator: NotIn
          values: [ kube-system, kueue-system ]
    #internalCertManagement:
    #  enable: false
    #  webhookServiceName: ""
    #  webhookSecretName: ""
    integrations:
      frameworks:
      - "batch/job"
      - "kubeflow.org/mpijob"
      - "ray.io/rayjob"
      - "ray.io/raycluster"
      - "jobset.x-k8s.io/jobset"
      - "kubeflow.org/paddlejob"
      - "kubeflow.org/pytorchjob"
      - "kubeflow.org/tfjob"
      - "kubeflow.org/xgboostjob"
      - "kubeflow.org/jaxjob"
      - "workload.codeflare.dev/appwrapper"
      - "pod"
    #  - "deployment" # requires enabling pod integration
    #  - "statefulset" # requires enabling pod integration
    #  - "leaderworkerset.x-k8s.io/leaderworkerset" # requires enabling pod integration
    #  externalFrameworks:
    #  - "Foo.v1.example.com"
    #fairSharing:
    #  enable: true
    #  preemptionStrategies: [LessThanOrEqualToFinalShare, LessThanInitialShare]
    #admissionFairSharing:
    #  usageHalfLifeTime: "168h" # 7 days
    #  usageSamplingInterval: "5m"
    #  resourceWeights: # optional, defaults to 1 for all resources if not specified
    #    cpu: 0    # if you want to completely ignore cpu usage
    #    memory: 0 # ignore completely memory usage
    #    example.com/gpu: 100 # and you care only about GPUs usage
    #resources:
    #  excludeResourcePrefixes: []
    #  transformations:
    #  - input: nvidia.com/mig-4g.5gb
    #    strategy: Replace | Retain
    #    outputs:
    #      example.com/accelerator-memory: 5Gi
    #      example.com/accelerator-gpc: 4
    #objectRetentionPolicies:
    #  workloads:
    #    afterFinished: null # null indicates infinite retention, 0s means no retention at all
    #    afterDeactivatedByKueue: null # null indicates infinite retention, 0s means no retention at all

  3. 在 Cloud Shell 中執行下列指令，安裝 Kueue：

    cd ${EXAMPLE_HOME}
kubectl kustomize kueue |kubectl apply --server-side --filename=-

    等待 Kueue Pod 準備就緒：

    watch kubectl --namespace=kueue-system get pods

    輸出內容應如下所示：

    NAME                                        READY   STATUS    RESTARTS   AGE
kueue-controller-manager-bdc956fc4-vhcmx    1/1     Running   0          3m15s

  4. 在 /workloads 目錄中，查看 flavors.yaml、cluster-queue.yaml 和 local-queue.yaml 檔案。這些資訊清單會指定 Kueue 管理資源配額的方式：

    ResourceFlavor

    這個資訊清單會在 Kueue 中定義預設的ResourceFlavor，用於資源管理。

    apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: default-flavor

    ClusterQueue

    這份資訊清單會設定 Kueue ClusterQueue，並為 CPU、記憶體和 GPU 設定資源限制。

    本教學課程使用的節點會附加兩個 Nvidia L4 GPU，對應的節點類型為 g2-standard-24，提供 24 個 vCPU 和 96 GB RAM。範例程式碼會說明如何將工作負載的資源用量限制在最多六個 GPU。

    ClusterQueue 設定中的 preemption 欄位會參照 PriorityClass，判斷資源不足時可先占哪些 Pod。

    apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  namespaceSelector: {} # match all.
  preemption:
    reclaimWithinCohort: LowerPriority
    withinClusterQueue: LowerPriority
  resourceGroups:
  - coveredResources: [ "cpu", "memory", "nvidia.com/gpu", "ephemeral-storage" ]
    flavors:
    - name: default-flavor
      resources:
      - name: "cpu"
        nominalQuota: 72
      - name: "memory"
        nominalQuota: 288Gi
      - name: "nvidia.com/gpu"
        nominalQuota: 6
      - name: "ephemeral-storage"
        nominalQuota: 200Gi

    LocalQueue

    這份資訊清單會在 llm 命名空間中建立名為 lq 的 Kueue LocalQueue。

    apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: llm # LocalQueue under llm namespace 
  name: lq
spec:
  clusterQueue: cluster-queue # Point to the ClusterQueue

  5. 查看 default-priorityclass.yaml、low-priorityclass.yaml 和 high-priorityclass.yaml 檔案。這些資訊清單會定義 Kubernetes 排程的 PriorityClass 物件。

    預設優先順序

    apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
  name: default-priority-nonpreempting
value: 10
preemptionPolicy: Never
globalDefault: true
description: "This priority class will not cause other pods to be preempted."

    低優先順序

    apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
  name: low-priority-preempting
value: 20
preemptionPolicy: PreemptLowerPriority
globalDefault: false
description: "This priority class will cause pods with lower priority to be preempted."

    高優先順序

    apiVersion: scheduling.k8s.io/v1
kind: PriorityClass
metadata:
  name: high-priority-preempting
value: 30
preemptionPolicy: PreemptLowerPriority
globalDefault: false
description: "This high priority class will cause other pods to be preempted."

  6. 執行下列指令來套用對應的資訊清單，藉此建立 Kueue 和 Kubernetes 物件。

    cd ${EXAMPLE_HOME}/workloads
kubectl apply --filename=flavors.yaml
kubectl apply --filename=default-priorityclass.yaml
kubectl apply --filename=high-priorityclass.yaml
kubectl apply --filename=low-priorityclass.yaml
kubectl apply --filename=cluster-queue.yaml
kubectl apply --filename=local-queue.yaml --namespace=llm

部署 TGI 推論伺服器

在本節中，您會部署 TGI 容器，提供 Gemma 2 模型。

  1. 在 /workloads 目錄中，查看 tgi-gemma-2-9b-it-hp.yaml 檔案。這份資訊清單定義了 Kubernetes Deployment，用於部署 TGI 服務執行階段和 gemma-2-9B-it 模型。Deployment 是 Kubernetes API 物件，可讓您執行多個 Pod 副本，並將這些副本分散到叢集中的節點。

    部署作業會優先處理推論工作，並為模型使用兩個 GPU。它會透過設定 NUM_SHARD 環境變數，使用張量平行化，將模型放入 GPU 記憶體。

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: tgi-gemma-deployment
  labels:
    app: gemma-server
spec:
  replicas: 1
  selector:
    matchLabels:
      app: gemma-server
  template:
    metadata:
      labels:
        app: gemma-server
        ai.gke.io/model: gemma-2-9b-it
        ai.gke.io/inference-server: text-generation-inference
        examples.ai.gke.io/source: user-guide
        kueue.x-k8s.io/queue-name: lq
    spec:
      priorityClassName: high-priority-preempting
      containers:
      - name: inference-server
        image: us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310
        resources:
          requests:
            cpu: "4"
            memory: "30Gi"
            ephemeral-storage: "30Gi"
            nvidia.com/gpu: "2"
          limits:
            cpu: "4"
            memory: "30Gi"
            ephemeral-storage: "30Gi"
            nvidia.com/gpu: "2"
        env:
        - name: AIP_HTTP_PORT
          value: '8000'
        - name: NUM_SHARD
          value: '2'
        - name: MODEL_ID
          value: google/gemma-2-9b-it
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        volumeMounts:
        - mountPath: /dev/shm
          name: dshm
      volumes:
      - name: dshm
        emptyDir:
          medium: Memory
      nodeSelector:
        cloud.google.com/gke-accelerator: "nvidia-l4"
---
apiVersion: v1
kind: Service
metadata:
  name: llm-service
spec:
  selector:
    app: gemma-server
  type: ClusterIP
  ports:
  - protocol: TCP
    port: 8000
    targetPort: 8000

  2. 執行下列指令來套用資訊清單：

    kubectl apply --filename=tgi-gemma-2-9b-it-hp.yaml --namespace=llm

    部署作業需要幾分鐘才能完成。

  3. 如要檢查 GKE 是否成功建立 Deployment，請執行下列指令：

    kubectl --namespace=llm get deployment

    輸出內容應如下所示：

    NAME                   READY   UP-TO-DATE   AVAILABLE   AGE
tgi-gemma-deployment   1/1     1            1           5m13s

驗證 Kueue 配額管理

在本節中，您將確認 Kueue 是否正確強制執行 Deployment 的 GPU 配額。

  1. 如要確認 Kueue 是否知道您的 Deployment，請執行下列指令來擷取 Workload 物件的狀態：

    kubectl --namespace=llm get workloads

    輸出內容應如下所示：

    NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
pod-tgi-gemma-deployment-6bf9ffdc9b-zcfrh-84f19   lq      cluster-queue   True                  8m23s

  2. 如要測試覆寫配額限制，請將 Deployment 擴展至四個副本：

    kubectl scale --replicas=4 deployment/tgi-gemma-deployment --namespace=llm

  3. 執行下列指令，查看 GKE 部署的副本數量：

    kubectl get workloads --namespace=llm

    輸出內容應如下所示：

    NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
pod-tgi-gemma-deployment-6cb95cc7f5-5thgr-3f7d4   lq      cluster-queue   True                  14s
pod-tgi-gemma-deployment-6cb95cc7f5-cbxg2-d9fe7   lq      cluster-queue   True                  5m41s
pod-tgi-gemma-deployment-6cb95cc7f5-tznkl-80f6b   lq                                            13s
pod-tgi-gemma-deployment-6cb95cc7f5-wd4q9-e4302   lq      cluster-queue   True                  13s

    輸出內容顯示，由於 Kueue 強制執行的資源配額，只有三個 Pod 獲得許可。

  4. 執行下列指令，顯示 llm 命名空間中的 Pod：

    kubectl get pod --namespace=llm

    輸出內容應如下所示：

    NAME                                    READY   STATUS            RESTARTS   AGE
tgi-gemma-deployment-7649884d64-6j256   1/1     Running           0          4m45s
tgi-gemma-deployment-7649884d64-drpvc   0/1     SchedulingGated   0          7s
tgi-gemma-deployment-7649884d64-thdkq   0/1     Pending           0          7s
tgi-gemma-deployment-7649884d64-znvpb   0/1     Pending           0          7s

  5. 現在，將 Deployment 縮減回 1。部署微調作業前，請務必完成這個步驟，否則系統會優先處理推論作業，導致微調作業無法順利進行。

    kubectl scale --replicas=1 deployment/tgi-gemma-deployment --namespace=llm

行為說明

由於您在 ClusterQueue 設定中設定了 GPU 配額限制，因此即使擴充至四個副本，擴充範例也只會產生三個副本。ClusterQueue 的 spec.resourceGroups 區段會為 nvidia.com/gpu 定義「6」的 nominalQuota。Deployment 會指定每個 Pod 需要「2」個 GPU。因此，ClusterQueue 一次最多只能容納三個 Deployment 副本 (因為 3 個副本 * 每個副本 2 個 GPU = 6 個 GPU，這是配額總數)。

當您嘗試擴充至四個副本時，Kueue 會發現這項動作會超出 GPU 配額，因此會防止排定第四個副本。第四個 Pod 的 SchedulingGated 狀態會指出這個情況。這個行為顯示 Kueue 的資源配額強制執行功能。

部署訓練工作

在本節中，您會為 Gemma 2 模型部署優先順序較低的微調工作，這項工作需要兩個 Pod 中的四個 GPU。Kubernetes 中的 Job 控制器會建立一或多個 Pod，並確保這些 Pod 成功執行特定工作。

這項 Job 會使用 ClusterQueue 中剩餘的 GPU 配額。這項作業會使用預先建構的映像檔，並儲存檢查點，以便從中繼結果重新啟動。

微調工作會使用 b-mc2/sql-create-context 資料集。您可以在存放區中找到微調工作的來源。

  1. 查看 fine-tune-l4.yaml 檔案。這個資訊清單定義了微調工作。

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc-l4
spec:
  clusterIP: None # clusterIP must be None to create a headless service
  selector:
    job-name: finetune-gemma-l4 # must match Job name
---
apiVersion: batch/v1
kind: Job
metadata:
  name: finetune-gemma-l4
  labels:
    kueue.x-k8s.io/queue-name: lq
spec:
  backoffLimit: 4
  completions: 2
  parallelism: 2
  completionMode: Indexed
  suspend: true # Set to true to allow Kueue to control the Job when it starts
  template:
    metadata:
      labels:
        app: finetune-job
      annotations:
        gke-gcsfuse/volumes: "true"
        gke-gcsfuse/memory-limit: "35Gi"
    spec:
      priorityClassName: low-priority-preempting
      containers:
      - name: gpu-job
        imagePullPolicy: Always
        image: us-docker.pkg.dev/google-samples/containers/gke/gemma-fine-tuning:v1.0.0
        ports:
        - containerPort: 29500
        resources:
          requests:
            nvidia.com/gpu: "2"
          limits:
            nvidia.com/gpu: "2"
        command:
        - bash
        - -c
        - |
          accelerate launch \
          --config_file fsdp_config.yaml \
          --debug \
          --main_process_ip finetune-gemma-l4-0.headless-svc-l4 \
          --main_process_port 29500 \
          --machine_rank ${JOB_COMPLETION_INDEX} \
          --num_processes 4 \
          --num_machines 2 \
          fine_tune.py
        env:
        - name: "EXPERIMENT"
          value: "finetune-experiment"
        - name: MODEL_NAME
          value: "google/gemma-2-2b"
        - name: NEW_MODEL
          value: "gemma-ft"
        - name: MODEL_PATH
          value: "/model-data/model-gemma2/experiment"
        - name: DATASET_NAME
          value: "b-mc2/sql-create-context"
        - name: DATASET_LIMIT
          value: "5000"
        - name: EPOCHS
          value: "1"
        - name: GRADIENT_ACCUMULATION_STEPS
          value: "2"
        - name: CHECKPOINT_SAVE_STEPS
          value: "10"
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        volumeMounts:
        - mountPath: /dev/shm
          name: dshm
        - name: gcs-fuse-csi-ephemeral
          mountPath: /model-data
          readOnly: false
      nodeSelector:
        cloud.google.com/gke-accelerator: nvidia-l4
      restartPolicy: OnFailure
      serviceAccountName: default
      subdomain: headless-svc-l4
      terminationGracePeriodSeconds: 60
      volumes:
      - name: dshm
        emptyDir:
          medium: Memory
      - name: gcs-fuse-csi-ephemeral
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: <MODEL_BUCKET>
            mountOptions: "implicit-dirs"
            gcsfuseLoggingSeverity: warning

  2. 套用資訊清單，建立微調作業：

    cd ${EXAMPLE_HOME}/workloads

sed -e "s/<MODEL_BUCKET>/$MODEL_BUCKET/g" \
    -e "s/<PROJECT_ID>/$PROJECT_ID/g" \
    -e "s/<REGION>/$REGION/g" \
    fine-tune-l4.yaml |kubectl apply --filename=- --namespace=llm

  3. 確認部署作業正在執行。如要檢查 Workload 物件的狀態，請執行下列指令：

    kubectl get workloads --namespace=llm

    輸出內容應如下所示：

    NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
job-finetune-gemma-l4-3316f                       lq      cluster-queue   True                  29m
pod-tgi-gemma-deployment-6cb95cc7f5-cbxg2-d9fe7   lq      cluster-queue   True                  68m

    接著執行下列指令，查看 llm 命名空間中的 Pod：

    kubectl get pod --namespace=llm

    輸出內容應如下所示：

    NAME                                    READY   STATUS    RESTARTS   AGE
finetune-gemma-l4-0-vcxpz               2/2     Running   0          31m
finetune-gemma-l4-1-9ppt9               2/2     Running   0          31m
tgi-gemma-deployment-6cb95cc7f5-cbxg2   1/1     Running   0          70m

    輸出內容顯示 Kueue 允許微調工作和推論伺服器 Pod 執行，並根據您指定的配額限制預留正確的資源。

  4. 查看輸出記錄，確認微調工作是否將檢查點儲存至 Cloud Storage 值區。微調作業約 10 分鐘後，就會開始儲存第一個檢查點。

    kubectl logs --namespace=llm --follow --selector=app=finetune-job

    第一個儲存的檢查點輸出內容如下：

    {"name": "finetune", "thread": 133763559483200, "threadName": "MainThread", "processName": "MainProcess", "process": 33, "message": "Fine tuning started", "timestamp": 1731002351.0016131, "level": "INFO", "runtime": 451579.89835739136}
…
{"name": "accelerate.utils.fsdp_utils", "thread": 136658669348672, "threadName": "MainThread", "processName": "MainProcess", "process": 32, "message": "Saving model to /model-data/model-gemma2/experiment/checkpoint-10/pytorch_model_fsdp_0", "timestamp": 1731002386.1763802, "level": "INFO", "runtime": 486753.8924217224}

在混合工作負載中測試 Kueue 先占和動態分配功能

在本節中，您將模擬推論伺服器負載增加的情境，因此需要擴充資源。這個情境示範了 Kueue 如何在資源受限時，暫停並搶占低優先順序的微調工作，優先處理高優先順序的推論伺服器。

  1. 執行下列指令，將推論伺服器的副本數擴充至兩個：

    kubectl scale --replicas=2 deployment/tgi-gemma-deployment --namespace=llm

  2. 檢查 Workload 物件的狀態：

    kubectl get workloads --namespace=llm

    輸出看起來類似以下內容：

    NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
job-finetune-gemma-l4-3316f                       lq                      False                 32m
pod-tgi-gemma-deployment-6cb95cc7f5-cbxg2-d9fe7   lq      cluster-queue   True                  70m
pod-tgi-gemma-deployment-6cb95cc7f5-p49sh-167de   lq      cluster-queue   True                  14s

    輸出內容顯示，由於增加的推論伺服器副本正在使用可用的 GPU 配額，因此微調作業不再獲准。

  3. 檢查微調工作的狀態：

    kubectl get job --namespace=llm

    輸出內容如下所示，表示微調作業狀態現已暫停：

    NAME                STATUS      COMPLETIONS   DURATION   AGE
finetune-gemma-l4   Suspended   0/2                      33m

  4. 執行下列指令來檢查 Pod：

    kubectl get pod --namespace=llm

    輸出內容類似於下列項目，表示 Kueue 終止了微調作業 Pod，以便為優先順序較高的推論伺服器 Deployment 釋出資源。

    NAME                                    READY   STATUS              RESTARTS   AGE
tgi-gemma-deployment-6cb95cc7f5-cbxg2   1/1     Running             0          72m
tgi-gemma-deployment-6cb95cc7f5-p49sh   0/1     ContainerCreating   0          91s

  5. 接著，測試推論伺服器負載減少，且 Pod 縮減規模的情況。執行下列指令：

    kubectl scale --replicas=1 deployment/tgi-gemma-deployment --namespace=llm

    執行下列指令來顯示 Workload 物件：

    kubectl get workloads --namespace=llm

    輸出內容如下所示，表示其中一個推論伺服器 Deployment 已終止，且微調作業已重新加入。

    NAME                                              QUEUE   RESERVED IN     ADMITTED   FINISHED   AGE
job-finetune-gemma-l4-3316f                       lq      cluster-queue   True                  37m
pod-tgi-gemma-deployment-6cb95cc7f5-cbxg2-d9fe7   lq      cluster-queue   True                  75m

  6. 執行下列指令來顯示 Jobs：

    kubectl get job --namespace=llm

    輸出內容會與下列內容類似，表示微調作業正在再次執行，並從最新的可用檢查點繼續執行。

    NAME                STATUS    COMPLETIONS   DURATION   AGE
finetune-gemma-l4   Running   0/2           2m11s      38m