混合 AI/ML トレーニング ワークロードと推論ワークロードの GKE リソース使用率を最適化する


このチュートリアルでは、単一の Google Kubernetes Engine(GKE)クラスタ内のトレーニング ワークロードと推論ワークロード間でアクセラレータ リソースを効率的に共有する方法について説明します。混合ワークロードを単一のクラスタに分散することで、リソース使用率が向上し、クラスタ管理が簡素化され、アクセラレータ数の制限による問題が軽減され、全体的な費用対効果が向上します。

このチュートリアルでは、推論用の Gemma 2 大規模言語モデル(LLM)と Hugging Face TGI(Text Generation Interface)サービング フレームワークを使用して、優先度の高いサービング Deployment を作成し、優先度の低い LLM ファインチューニング Job を作成します。どちらのワークロードも、NVIDIA L4 GPU を使用する単一クラスタで実行されます。オープンソースの Kubernetes ネイティブの Job キューイング システムである Kueue を使用して、ワークロードを管理してスケジュールします。Kueue を使用すると、サービング タスクの優先度を設定したり、優先度の低いトレーニング Job をプリエンプトしたりして、リソース使用率を最適化できます。サービング デマンドの減少に伴い、解放されたアクセラレータを再割り当てして、トレーニング ジョブを再開します。Kueue と優先度クラスを使用して、プロセス全体でリソース割り当てを管理します。

このチュートリアルは、GKE クラスタで ML モデルをトレーニングしてホストし、特に限られた数のアクセラレータを使用する場合に、費用と管理のオーバーヘッドを削減する必要がある ML エンジニア、プラットフォーム管理者、オペレーター、データおよび AI スペシャリストを対象としています。コンテンツで参照する一般的なロールとタスク例の詳細については、一般的な GKE Enterprise ユーザーロールとタスクをご覧ください。 Google Cloud

このページを読む前に、次のことをよく理解しておいてください。

目標

このガイドを終えると、次の手順を行えるようになります。

  • 優先度の高いサービング Deployment を構成します。
  • 優先度の低いトレーニング ジョブを設定します。
  • 需要の変化に対応するためにプリエンプション戦略を実装する。
  • Kueue を使用して、トレーニング タスクとサービング タスク間のリソース割り当てを管理します。

始める前に

  • Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required APIs.

    Enable the APIs

  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  • Make sure that billing is enabled for your Google Cloud project.

  • Enable the required APIs.

    Enable the APIs

  • Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin

    Check for the roles

    1. In the Google Cloud console, go to the IAM page.

      Go to IAM
    2. Select the project.
    3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

    4. For all rows that specify or include you, check the Role column to see whether the list of roles includes the required roles.

    Grant the roles

    1. In the Google Cloud console, go to the IAM page.

      [IAM] に移動
    2. プロジェクトを選択します。
    3. [ アクセスを許可] をクリックします。
    4. [新しいプリンシパル] フィールドに、ユーザー ID を入力します。 これは通常、Google アカウントのメールアドレスです。

    5. [ロールを選択] リストでロールを選択します。
    6. 追加のロールを付与するには、 [別のロールを追加] をクリックして各ロールを追加します。
    7. [保存] をクリックします。

環境を準備する

このセクションでは、推論ワークロードとトレーニング ワークロード用の TGI とモデルをデプロイするために必要なリソースをプロビジョニングします。

モデルへのアクセス権を取得する

GKE にデプロイするために Gemma モデルへのアクセス権を取得するには、まずライセンス同意契約に署名してから、Hugging Face のアクセス トークンを生成する必要があります。

  1. ライセンス同意契約に署名します。モデルの同意ページにアクセスし、Hugging Face アカウントを使用して同意を確認し、モデルの利用規約に同意します。
  2. アクセス トークンを生成する。Hugging Face からモデルにアクセスするには、Hugging Face トークンが必要です。トークンをまだ生成していない場合は、次の手順に沿って生成します。

    1. [Your Profile] > [Settings] > [Access Tokens] の順にクリックします。
    2. [New Token] を選択します。
    3. 任意の名前と、少なくとも Read ロールを指定します。
    4. [Generate a token] を選択します。
    5. トークンをクリップボードにコピーします。

Cloud Shell を起動する

このチュートリアルでは、Cloud Shell を使用してGoogle Cloudでホストされているリソースを管理します。Cloud Shell には、このチュートリアルに必要な kubectl gcloud CLITerraform などのソフトウェアがプリインストールされています。

Cloud Shell を使用して環境を設定するには、次の操作を行います。

  1. Google Cloud コンソールで、Google Cloud コンソールCloud Shell 有効化アイコン [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 クラスタまたは Standard クラスタを使用できます。フルマネージドの Kubernetes エクスペリエンスを実現するには、Autopilot クラスタを使用することをおすすめします。ワークロードに最適な 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} \
        --region=${REGION} \
        --release-channel=rapid
    
  3. ファインチューニング ジョブ用の Cloud Storage バケットを作成します。

    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 \
        --region=$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 の使用の制限事項をご覧ください。たとえば、このチュートリアルでは、ノードごとに 2 つの GPU が接続された g2-standard-24 を使用します。使用可能な GPU の最新のリストについては、コンピューティング ワークロード用の GPU をご覧ください。
    • GPU_POOL_ACCELERATOR_TYPE: 選択したリージョンでサポートされているアクセラレータのタイプ。たとえば、このチュートリアルでは nvidia-l4 を使用します。使用可能な GPU の最新のリストについては、コンピューティング ワークロード用の GPU をご覧ください。
    • MODEL_BUCKET: トレーニング済みモデルの重みを保存する Cloud Storage バケット。
  2. Standard クラスタを作成します。

    gcloud container clusters create ${CLUSTER_NAME} \
        --project=${PROJECT_ID} \
        --region=${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 バケットを作成します。

    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 \
        --region=$REGION \
        --project=$PROJECT_ID
    
  7. Deployment の名前空間を作成します。Cloud Shell で、次のコマンドを実行します。

    kubectl create ns llm
    

Hugging Face の認証情報用の Kubernetes Secret を作成する

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 を制御するには、Kustomize を使用してカスタム構成を適用し、v1/pod 統合を有効にして、サーバー Pod に "kueue-job: true" というラベルを付けます。

  1. /kueue ディレクトリで、kustomization.yaml のコードを確認します。このマニフェストは、カスタム構成で Kueue リソース マネージャーをインストールします。

    apiVersion: kustomize.config.k8s.io/v1beta1
    kind: Kustomization
    resources:
    - https://github.com/kubernetes-sigs/kueue/releases/download/v0.10.0/manifests.yaml
    patches:
    - path: patch.yaml
      target:
        version: v1
        kind: ConfigMap
        name: kueue-manager-config
    
  2. /kueue ディレクトリで、patch.yaml のコードを確認します。この ConfigMap は、"kueue-job: true" ラベルを持つ Pod を管理するように Kueue をカスタマイズします。

    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
        #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/mxjob"
          - "kubeflow.org/paddlejob"
          - "kubeflow.org/pytorchjob"
          - "kubeflow.org/tfjob"
          - "kubeflow.org/xgboostjob"
          - "pod"
        #  externalFrameworks:
        #  - "Foo.v1.example.com"
          podOptions:
            # You can change namespaceSelector to define in which 
            # namespaces kueue will manage the pods.
            namespaceSelector:
              matchExpressions:
              - key: kubernetes.io/metadata.name
                operator: NotIn
                values: [ kube-system, kueue-system ]
            # Kueue uses podSelector to manage pods with particular 
            # labels. The default podSelector will match all the pods. 
            podSelector:
              matchExpressions:
              - key: kueue-job
                operator: In
                values: [ "true", "True", "yes" ]
    
  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    2/2     Running   0          3m15s
    
  4. /workloads ディレクトリで、flavors.yamlcluster-queue.yamllocal-queue.yaml ファイルを表示します。これらのマニフェストには、Kueue がリソース割り当てを管理する方法を指定します。

    ResourceFlavor

    このマニフェストは、リソース管理のために Kueue のデフォルトの ResourceFlavor を定義します。

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

    ClusterQueue

    このマニフェストは、CPU、メモリ、GPU のリソース上限を持つ Kueue の ClusterQueue を設定します。

    このチュートリアルでは、2 つの Nvidia L4 GPU が接続されたノードを使用し、対応するノードタイプは g2-standard-24 で、24 個の vCPU と 96 GB の RAM を提供します。次のコードサンプルは、ワークロードのリソース使用量を最大 6 個の GPU に制限する方法を示しています。

    ClusterQueue 構成の preemption フィールドは、PriorityClasses を参照して、リソースが不足している場合にプリエンプトできる 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 Namespace に 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.yamllow-priorityclass.yamlhigh-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 推論サーバーをデプロイする

このセクションでは、Gemma 2 モデルを提供する TGI コンテナをデプロイします。

  1. /workloads ディレクトリで tgi-gemma-2-9b-it-hp.yaml ファイルを表示します。このマニフェストは、TGI サービング ランタイムと gemma-2-9B-it モデルをデプロイする Kubernetes Deployment を定義します。Deployment は、クラスタ内のノードに分散された Pod の複数のレプリカを実行できる Kubernetes API オブジェクトです。

    Deployment は推論タスクを優先し、モデルに 2 つの 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
            kueue-job: "true"
        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 を 4 つのレプリカにスケールします。

    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 が適用するリソース割り当てにより、3 つの Pod のみが許可されていることが示されます。

  4. 次のコマンドを実行して、llm Namespace 内の 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 割り当て上限により、4 つにスケーリングしてもレプリカは 3 つしか作成されません。ClusterQueue の spec.resourceGroups セクションでは、nvidia.com/gpu の nominalQuota を「6」と定義しています。Deployment は、各 Pod に「2」個の GPU が必要であることを指定します。したがって、ClusterQueue は一度に Deployment のレプリカを最大 3 つしか処理できません(3 つのレプリカ * レプリカあたり 2 つの GPU = 6 つの GPU が合計割り当てであるため)。

4 つのレプリカにスケーリングしようとすると、このアクションが GPU 割り当てを超えることを Kueue が認識し、4 番目のレプリカがスケジュールされなくなります。これは、4 番目の Pod の SchedulingGated ステータスによって示されます。この動作は、Kueue のリソース割り当ての適用を示しています。

トレーニング Job をデプロイする

このセクションでは、2 つの Pod に 4 つの GPU を必要とする Gemma 2 モデルの優先度の低いファインチューニング ジョブをデプロイします。Kubernetes の Job コントローラは、1 つ以上の Pod を作成し、特定のタスクが正常に実行されるようにします。

この Job は、ClusterQueue の残りの GPU 割り当てを使用します。Job はビルド済みイメージを使用し、中間結果から再起動できるようにチェックポイントを保存します。

ファインチューニング ジョブは b-mc2/sql-create-context データセットを使用します。ファインチューニング ジョブのソースは、リポジトリにあります。

  1. fine-tune-l4.yaml ファイルを表示します。このマニフェストは、ファインチューニング Job を定義します。

    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. Deployment が実行されていることを確認します。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 Namespace の 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 がファインチューニング Job と推論サーバー Pod の両方の実行を許可し、指定された割り当て上限に基づいて適切なリソースを予約していることがわかります。

  4. 出力ログを表示して、ファインチューニング ジョブがチェックポイントを Cloud Storage バケットに保存していることを確認します。ファインチューニング Job が最初のチェックポイントの保存を開始するまでに 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. 次のコマンドを実行して、推論サーバーのレプリカを 2 つにスケールします。

    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. ファインチューニング Job のステータスを確認します。

    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 の 1 つが終了し、ファインチューニング ジョブが再承認されたことを示します。

    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. 次のコマンドを実行して Job を表示します。

    kubectl get job --namespace=llm
    

    出力は次のようになります。これは、ファインチューニング ジョブが再び実行され、利用可能な最新のチェックポイントから再開されたことを示します。

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

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

デプロイされたリソースを削除する

このガイドで作成したリソースについて、 Google Cloud アカウントに課金されないようにするには、次のコマンドを実行します。

gcloud storage rm --recursive gs://${MODEL_BUCKET}
gcloud container clusters delete ${CLUSTER_NAME} --location ${REGION}

次のステップ