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

Standard Autopilot

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

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

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

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

目標

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

優先度の高いサービング Deployment を構成します。
優先度の低いトレーニング Jobs を設定します。
需要の変化に対応するためにプリエンプション戦略を実装します。
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

In the Google Cloud console, go to the IAM page.
Go to IAM
Select the project.
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.
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

In the Google Cloud console, go to the IAM page.
IAM に移動
プロジェクトを選択します。
[ アクセスを許可] をクリックします。
[新しいプリンシパル] フィールドに、ユーザー ID を入力します。これは通常、Google アカウントのメールアドレスです。
[ロールを選択] リストでロールを選択します。
追加のロールを付与するには、 [別のロールを追加] をクリックして各ロールを追加します。
[保存] をクリックします。

Hugging Face アカウントを作成します（まだ作成していない場合）。
L4 GPU 用にプロジェクトに十分な割り当てがあることを確認します。詳細については、GPU についてと数量に基づく割り当てをご覧ください。

環境を準備する

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

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

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

ライセンス同意契約に署名する。モデルの同意ページにアクセスし、Hugging Face アカウントを使用して同意を確認して、モデルの利用規約に同意します。
アクセストークンを生成する。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 CLI、Terraform などのソフトウェアがプリインストールされています。

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

Google Cloud コンソールで（Cloud Shell をアクティブにする）をクリックして、Google Cloud コンソールで Cloud Shell セッションを起動します。これにより、 Google Cloud コンソールの下部ペインでセッションが起動します。
デフォルトの環境変数を設定します。
```
gcloud config set project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)
```
PROJECT_ID は、実際の Google Cloudプロジェクト ID に置き換えます。

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

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 バケットを表します。

Autopilot クラスタを作成します。

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

ファインチューニング Job 用の Cloud Storage バケットを作成します。

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

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

クラスタの認証情報を取得するには、次のコマンドを実行します。

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

Deployment の名前空間を作成します。Cloud Shell で、次のコマンドを実行します。
```
kubectl create ns llm
```

標準

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 バケット。

Standard クラスタを作成します。

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

推論とファインチューニングのワークロード用の 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

ファインチューニング Job 用の Cloud Storage バケットを作成します。

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

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

クラスタの認証情報を取得するには、次のコマンドを実行します。

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

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 のプリエンプション機能は、リソースが不足しているときに優先度の低いトレーニング Jobs を一時停止または強制排除することで、優先度の高いサービングワークロードに常に必要なリソースを確保します。

Kueue で推論サーバーの Deployment を制御するには、pod インテグレーションを有効にして、kube-system Namespace と kueue-system Namespace を除外するように managedJobsNamespaceSelector を構成します。

/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

/kueue ディレクトリで、patch.yaml のコードを確認します。この ConfigMap は、kube-system Namespace と kueue-system Namespace の 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
    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

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

/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
このマニフェストは、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 名前空間に 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

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."

次のコマンドを実行して対応するマニフェストを適用し、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 コンテナをデプロイします。

/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
    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

次のコマンドを実行してマニフェストを適用します。
```
kubectl apply --filename=tgi-gemma-2-9b-it-hp.yaml --namespace=llm
```
デプロイオペレーションが完了するまでに数分かかります。

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 割り当てを正しく適用していることを確認します。

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

割り当て上限のオーバーライドをテストするには、Deployment を 4 つのレプリカにスケーリングします。
```
kubectl scale --replicas=4 deployment/tgi-gemma-deployment --namespace=llm
```

次のコマンドを実行して、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 のみが許可されていることが示されます。

次のコマンドを実行して、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

Deployment を 1 にスケールダウンします。この手順は、ファインチューニング Job をデプロイする前に行う必要があります。この手順を実行しないと、推論 Job が優先されるため、Job は承認されません。
```
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 モデルの優先度の低いファインチューニング Job をデプロイします。Kubernetes の Job コントローラは、1 つ以上の Pod を作成し、特定のタスクが正常に実行されるようにします。

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

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

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

マニフェストを適用して、ファインチューニング Job を作成します。

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

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 名前空間の 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 の両方の実行を許可し、指定された割り当て上限に基づいて適切なリソースを予約していることがわかります。

出力ログを表示して、ファインチューニング Job がチェックポイントを 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 プリエンプションと動的割り当てをテストする

このセクションでは、推論サーバーの負荷が増加し、スケールアップが必要になるシナリオをシミュレートします。このシナリオでは、リソースが制約されているときに、優先度の低いファインチューニング Job を停止してプリエンプトすることで、Kueue が優先度の高い推論サーバーの優先度を上げる方法を示します。

次のコマンドを実行して、推論サーバーのレプリカを 2 つにスケールします。
```
kubectl scale --replicas=2 deployment/tgi-gemma-deployment --namespace=llm
```

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 割り当てを使用しているため、ファインチューニング Job が許可されなくなったことを示しています。

ファインチューニング Job のステータスを確認します。
```
kubectl get job --namespace=llm
```
出力は次のようになります。これは、ファインチューニング Job のステータスが一時停止されたことを示します。
```
NAME                STATUS      COMPLETIONS   DURATION   AGE
finetune-gemma-l4   Suspended   0/2                      33m
```

次のコマンドを実行して Pod を調べます。

kubectl get pod --namespace=llm

出力は次のようになります。これは、Kueue がファインチューニング Job 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

次に、推論サーバーの負荷が減少し、Pod がスケールダウンされるシナリオをテストします。次のコマンドを実行します。
```
kubectl scale --replicas=1 deployment/tgi-gemma-deployment --namespace=llm
```
次のコマンドを実行して、Workload オブジェクトを表示します。
```
kubectl get workloads --namespace=llm
```
出力は次のようになります。これは、推論サーバー Deployment の 1 つが終了し、ファインチューニング Job が再承認されたことを示します。
```
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
```
次のコマンドを実行して Job を表示します。
```
kubectl get job --namespace=llm
```
出力は次のようになります。これは、ファインチューニング Job が再び実行され、利用可能な最新のチェックポイントから再開されたことを示します。
```
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}

混合 AI/ML トレーニング ワークロードと推論ワークロードの GKE リソース使用率を最適化する コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

目標

始める前に

Check for the roles

Grant the roles

環境を準備する

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

Cloud Shell を起動する

GKE クラスタを作成する

Autopilot

標準

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

Kueue を構成する

ResourceFlavor

ClusterQueue

LocalQueue

デフォルトの優先値

低い優先度

優先度高

TGI 推論サーバーをデプロイする

Kueue の割り当て管理を確認する

動作の説明

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

混合ワークロードで Kueue プリエンプションと動的割り当てをテストする

クリーンアップ

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

次のステップ

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