このページは Cloud Translation API によって翻訳されました。

GKE で TPU Trillium と vLLM を使用して LLM をサービングする

このチュートリアルでは、vLLM サービングフレームワークを使用して、Google Kubernetes Engine（GKE）で Tensor Processing Unit（TPU）を使用して大規模言語モデル（LLM）をサービングする方法について説明します。このチュートリアルでは、Llama 3.1 70b をサービングし、TPU Trillium を使用します。また、vLLM サーバー指標を使用して水平 Pod 自動スケーリングを設定します。

このドキュメントは、AI / ML ワークロードをデプロイしてサービングする際に、マネージド Kubernetes での詳細な制御、スケーラビリティ、復元力、ポータビリティ、費用対効果が求められる場合の出発点として適しています。

背景

GKE で TPU Trillium を使用すると、効率的なスケーラビリティや高可用性をはじめとするマネージド Kubernetes のメリットをすべて活用し、プロダクションレディで堅牢なサービングソリューションを実装できます。このセクションでは、このガイドで使用されている重要なテクノロジーについて説明します。

TPU Trillium

TPU は、Google が独自に開発した特定用途向け集積回路（ASIC）です。TPU は、TensorFlow、PyTorch、JAX などのフレームワークを使用して構築された AI / ML モデルを高速化するために使用されます。このチュートリアルでは、Google の第 6 世代 TPU である TPU Trillium を使用します。

GKE で TPU を使用する前に、次の学習プログラムを完了することをおすすめします。

TPU Trillium のシステムアーキテクチャについて学習する。
GKE の TPU についてを確認する。

vLLM

vLLM は、LLM のサービング用に高度に最適化されたオープンソースフレームワークです。vLLM は、次のような機能により TPU でのサービングスループットを向上させることができます。

PagedAttention による Transformer の実装の最適化
サービングスループットを全体的に向上させる連続的なバッチ処理
複数の TPU でのテンソル並列処理と分散サービング

詳細については、vLLM のドキュメントをご覧ください。

注: このチュートリアルでは、単一ホスト構成での vLLM のデプロイに焦点を当てています。これは、ct6e-standard-8t マシンタイプの Llama 3.1 70b など、単一の TPU スライスからサービングできるモデルに最適です。GKE で TPU を使用して vLLM を使用する場合、マルチホスト構成はサポートされていません。マルチホストのサポートがないため、複数のホストの集約されたメモリとコンピューティングを必要とする非常に大きなモデル（4,000 億以上のパラメータなど）のサービングに vLLM を使用することはできません。マルチホスト設定が必要な本番環境システムまたはモデルの場合、推奨されるパフォーマンス最適化ソリューションは、TPU 推論用の Google エンジンである JetStream を使用することです。マルチホストデプロイを開始するには、v6e TPU での JetStream MaxText 推論をご覧ください。

Cloud Storage FUSE

Cloud Storage FUSE は、オブジェクトストレージバケットに存在するモデルの重み付けで、GKE クラスタから Cloud Storage にアクセスできるようにします。このチュートリアルで作成された Cloud Storage バケットは最初は空になります。vLLM が起動すると、GKE は Hugging Face からモデルをダウンロードし、重みを Cloud Storage バケットのキャッシュに保存します。Pod の再起動またはデプロイのスケールアップ時に、後続のモデル読み込みでキャッシュに保存されたデータが Cloud Storage バケットからダウンロードされ、並列ダウンロードを利用してパフォーマンスが最適化されます。

詳細については、Cloud Storage FUSE CSI ドライバのドキュメントをご覧ください。

GKE クラスタを作成する

GKE Autopilot クラスタまたは GKE Standard クラスタの TPU で LLM を提供できます。フルマネージドの Kubernetes エクスペリエンスを実現するには、Autopilot クラスタを使用することをおすすめします。ワークロードに最適な GKE の運用モードを選択するには、GKE の運用モードを選択するをご覧ください。

Autopilot

GKE Autopilot クラスタを作成します。
```
gcloud container clusters create-auto ${CLUSTER_NAME} \
    --cluster-version=${CLUSTER_VERSION} \
    --location=${CONTROL_PLANE_LOCATION}
```
注: チュートリアルの実行中に Cloud Shell インスタンスが切断された場合は、上記の手順を繰り返してください。

Standard

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

gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${CONTROL_PLANE_LOCATION} \
    --node-locations=${ZONE} \
    --cluster-version=${CLUSTER_VERSION} \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --addons GcsFuseCsiDriver

TPU スライスノードプールを作成します。
```
gcloud container node-pools create tpunodepool \
    --location=${CONTROL_PLANE_LOCATION} \
    --node-locations=${ZONE} \
    --num-nodes=1 \
    --machine-type=ct6e-standard-8t \
    --cluster=${CLUSTER_NAME} \
    --enable-autoscaling --total-min-nodes=1 --total-max-nodes=2
```
GKE は、LLM 用に次のリソースを作成します。
- Workload Identity Federation for GKE を使用し、Cloud Storage FUSE CSI ドライバが有効になっている GKE Standard クラスタ。
- ct6e-standard-8t マシンタイプの TPU Trillium ノードプール。このノードプールには 1 つのノードと 8 個の TPU チップがあり、自動スケーリングが有効になっています。

クラスタと通信するように kubectl を構成する

クラスタと通信するように kubectl を構成するには、次のコマンドを実行します。

  gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${CONTROL_PLANE_LOCATION}

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

Namespace を作成します。default Namespace を使用している場合は、この手順をスキップできます。
```
kubectl create namespace ${NAMESPACE}
```

Hugging Face トークンを含む Kubernetes Secret を作成するには、次のコマンドを実行します。

kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --namespace ${NAMESPACE}

Cloud Storage バケットを作成する

Cloud Shell で、次のコマンドを実行します。

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

これにより、Hugging Face からダウンロードしたモデルファイルを格納する Cloud Storage バケットが作成されます。

バケットにアクセスする Kubernetes ServiceAccount を設定する

Kubernetes ServiceAccount を作成します。

kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}

Cloud Storage バケットにアクセスできるように、Kubernetes ServiceAccount に読み取り / 書き込みアクセス権を付与します。

gcloud storage buckets add-iam-policy-binding gs://${GSBUCKET} \
  --member "principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}" \
  --role "roles/storage.objectUser"

また、プロジェクト内のすべての Cloud Storage バケットに対する読み取り / 書き込みアクセス権を付与することもできます。
```
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
--member "principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/${NAMESPACE}/sa/${KSA_NAME}" \
--role "roles/storage.objectUser"
```
GKE は、LLM 用に次のリソースを作成します。
1. ダウンロードしたモデルとコンパイルキャッシュを保存する Cloud Storage バケット。Cloud Storage FUSE CSI ドライバがバケットのコンテンツを読み取ります。
2. ファイルキャッシュが有効になっているボリュームと、Cloud Storage FUSE の並列ダウンロード機能。
ベストプラクティス:
モデルコンテンツ（重み付けファイルなど）の予想サイズに応じて、tmpfs または Hyperdisk / Persistent Disk を基盤とするファイルキャッシュを使用します。このチュートリアルでは、RAM を基盤とする Cloud Storage FUSE ファイルキャッシュを使用します。

vLLM モデルサーバーをデプロイする

このチュートリアルでは、vLLM モデルサーバーをデプロイするために Kubernetes Deployment を使用します。Deployment は、クラスタ内のノードに分散された Pod の複数のレプリカを実行できる Kubernetes API オブジェクトです。

単一のレプリカを使用する次の Deployment マニフェスト（vllm-llama3-70b.yaml として保存）を調べます。


apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-tpu
spec:
  replicas: 1
  selector:
    matchLabels:
      app: vllm-tpu
  template:
    metadata:
      labels:
        app: vllm-tpu
      annotations:
        gke-gcsfuse/volumes: "true"
        gke-gcsfuse/cpu-limit: "0"
        gke-gcsfuse/memory-limit: "0"
        gke-gcsfuse/ephemeral-storage-limit: "0"
    spec:
      serviceAccountName: KSA_NAME
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 2x4
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
      containers:
      - name: vllm-tpu
        image: vllm/vllm-tpu:latest
        command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
        args:
        - --host=0.0.0.0
        - --port=8000
        - --tensor-parallel-size=8
        - --max-model-len=4096
        - --model=meta-llama/Llama-3.1-70B
        - --download-dir=/data
        - --max-num-batched-tokens=512
        - --max-num-seqs=128
        env: 
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        - name: VLLM_XLA_CACHE_PATH
          value: "/data"
        - name: VLLM_USE_V1
          value: "1"
        ports:
        - containerPort: 8000
        resources:
          limits:
            google.com/tpu: 8
        readinessProbe:
          tcpSocket:
            port: 8000
          initialDelaySeconds: 15
          periodSeconds: 10
        volumeMounts:
        - name: gcs-fuse-csi-ephemeral
          mountPath: /data
        - name: dshm
          mountPath: /dev/shm
      volumes:
      - name: gke-gcsfuse-cache
        emptyDir:
          medium: Memory
      - name: dshm
        emptyDir:
          medium: Memory
      - name: gcs-fuse-csi-ephemeral
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: GSBUCKET
            mountOptions: "implicit-dirs,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:100,file-cache:max-parallel-downloads:-1,file-cache:download-chunk-size-mb:10,file-cache:max-size-mb:-1"
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-service
spec:
  selector:
    app: vllm-tpu
  type: LoadBalancer	
  ports:
    - name: http
      protocol: TCP
      port: 8000  
      targetPort: 8000

Deployment を複数のレプリカにスケールアップすると、VLLM_XLA_CACHE_PATH への同時書き込みにより、RuntimeError: filesystem error: cannot create directories エラーが発生します。このエラーを防ぐには、次の 2 つの方法があります。

Deployment YAML から次のブロックを削除して、XLA キャッシュの場所を削除します。つまり、すべてのレプリカでキャッシュが再コンパイルされます。
```
- name: VLLM_XLA_CACHE_PATH
  value: "/data"
```
Deployment を 1 にスケーリングし、最初レプリカの準備が整い、XLA キャッシュに書き込まれるまで待ちます。次に、追加のレプリカにスケーリングします。これにより、残りのレプリカはキャッシュを読み取ることができます。書き込みは試行されません。

次のコマンドを実行してマニフェストを適用します。
```
kubectl apply -f vllm-llama3-70b.yaml -n ${NAMESPACE}
```

実行中のモデルサーバーのログを表示します。

kubectl logs -f -l app=vllm-tpu -n ${NAMESPACE}

出力は次のようになります。

INFO:     Started server process [1]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)

モデルをサービングする

VLLM サービスの外部 IP アドレスを取得するには、次のコマンドを実行します。

export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}' -n ${NAMESPACE})

curl を使用してモデルを操作します。

curl http://$vllm_service:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
    "model": "meta-llama/Llama-3.1-70B",
    "prompt": "San Francisco is a",
    "max_tokens": 7,
    "temperature": 0
}'

出力例を以下に示します。

{"id":"cmpl-6b4bb29482494ab88408d537da1e608f","object":"text_completion","created":1727822657,"model":"meta-llama/Llama-3-8B","choices":[{"index":0,"text":" top holiday destination featuring scenic beauty and","logprobs":null,"finish_reason":"length","stop_reason":null,"prompt_logprobs":null}],"usage":{"prompt_tokens":5,"total_tokens":12,"completion_tokens":7}}

カスタムオートスケーラーを設定する

このセクションでは、カスタム Prometheus 指標を使用して水平 Pod 自動スケーリングを設定します。vLLM サーバーから Google Cloud Managed Service for Prometheus の指標を使用します。

詳細については、Google Cloud Managed Service for Prometheus をご覧ください。これは GKE クラスタでデフォルトで有効になっています。

クラスタにカスタム指標の Stackdriver アダプタを設定します。

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml

カスタム指標の Stackdriver アダプタが使用するサービスアカウントに Monitoring 閲覧者のロールを追加します。
```
gcloud projects add-iam-policy-binding projects/${PROJECT_ID} \
    --role roles/monitoring.viewer \
    --member=principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/custom-metrics/sa/custom-metrics-stackdriver-adapter
```
注: GKE クラスタで使用するサービスアカウントにモニタリング指標の書き込みロールがあることを確認します。このチュートリアルでは、デフォルトの Compute Engine サービスアカウントを使用します。

次のマニフェストを vllm_pod_monitor.yaml として保存します。


apiVersion: monitoring.googleapis.com/v1
kind: PodMonitoring
metadata:
 name: vllm-pod-monitoring
spec:
 selector:
   matchLabels:
     app: vllm-tpu
 endpoints:
 - path: /metrics
   port: 8000
   interval: 15s

クラスタに適用します。

kubectl apply -f vllm_pod_monitor.yaml -n ${NAMESPACE}

vLLM エンドポイントに負荷を生成する

vLLM サーバーに負荷をかけて、GKE がカスタム vLLM 指標で自動スケーリングする方法を確認します。

bash スクリプト（load.sh）を実行して、N 個の並列リクエストを vLLM エンドポイントに送信します。

#!/bin/bash
N=PARALLEL_PROCESSES
export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}' -n ${NAMESPACE})
for i in $(seq 1 $N); do
  while true; do
    curl http://$vllm_service:8000/v1/completions -H "Content-Type: application/json" -d '{"model": "meta-llama/Llama-3.1-70B", "prompt": "Write a story about san francisco", "max_tokens": 1000, "temperature": 0}'
  done &  # Run in the background
done
wait

PARALLEL_PROCESSES は、実行する並列プロセスの数に置き換えます。

bash スクリプトを実行します。
```
chmod +x load.sh
nohup ./load.sh &
```

Google Cloud Managed Service for Prometheus が指標を取り込むことを確認する

Google Cloud Managed Service for Prometheus が指標をスクレイピングし、vLLM エンドポイントに負荷をかけると、Cloud Monitoring で指標を表示できます。

Google Cloud コンソールで、Metrics Explorer のページに移動します。

Metrics Explorer に移動
[< > PromQL] をクリックします。
次のクエリを入力して、トラフィック指標を確認します。
```
vllm:num_requests_waiting{cluster='CLUSTER_NAME'}
```

折れ線グラフに、時間の経過に伴う vLLM 指標（num_requests_waiting）が表示されます。vLLM 指標は、0（プリロード）から値（ポストロード）にスケールアップします。このグラフでは、vLLM 指標が Google Cloud Managed Service for Prometheus に取り込まれていることがわかります。次のグラフの例は、プリロード開始時の値が 0 で、1 分以内に最大の負荷に達した後、値が 400 近くに達していることを示しています。

vllm:num_requests_waiting

HorizontalPodAutoscaler 構成をデプロイする

自動スケーリングする指標を決定する場合は、vLLM TPU に次の指標を使用することをおすすめします。

num_requests_waiting: この指標は、モデルサーバーのキューで待機しているリクエストの数に関連しています。この数は、kv キャッシュがいっぱいになると著しく増加します。
gpu_cache_usage_perc: この指標は kv キャッシュの使用率に関連しており、モデルサーバーで特定の推論サイクルで処理されるリクエスト数に直接関連しています。この指標は GPU と TPU で同じように機能しますが、GPU の命名スキーマに関連付けられています。

スループットと費用を最適化する場合、また、モデルサーバーの最大スループットでレイテンシ目標を達成できる場合は、num_requests_waiting を使用することをおすすめします。

キューベースのスケーリングでは要件を満たせない、レイテンシの影響を受けやすいワークロードがある場合は、gpu_cache_usage_perc を使用することをおすすめします。

詳細については、TPU を使用して大規模言語モデル（LLM）推論ワークロードを自動スケーリングするためのベストプラクティスをご覧ください。

HPA 構成の averageValue ターゲットを選択する場合は、テストで決定する必要があります。この部分を最適化する方法については、GPU のコストを削減: GKE の推論ワークロード向けのスマートな自動スケーリングのブログ投稿をご覧ください。このブログ投稿で使用した profile-generator は vLLM TPU でも機能します。

次の手順では、num_requests_waiting 指標を使用して HPA 構成をデプロイします。デモ用に指標を低い値に設定し、HPA 構成で vLLM レプリカを 2 にスケーリングします。num_requests_waiting を使用して Horizontal Pod Autoscaler 構成をデプロイする手順は次のとおりです。

次のマニフェストを vllm-hpa.yaml として保存します。


apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
 name: vllm-hpa
spec:
 scaleTargetRef:
   apiVersion: apps/v1
   kind: Deployment
   name: vllm-tpu
 minReplicas: 1
 maxReplicas: 2
 metrics:
   - type: Pods
     pods:
       metric:
         name: prometheus.googleapis.com|vllm:num_requests_waiting|gauge
       target:
         type: AverageValue
         averageValue: 10

Google Cloud Managed Service for Prometheus の vLLM 指標は vllm:metric_name 形式に従います。

ベストプラクティス:

スループットをスケーリングするには num_requests_waiting を使用します。レイテンシの影響を受けやすい TPU のユースケースには gpu_cache_usage_perc を使用します。

HorizontalPodAutoscaler 構成をデプロイします。
```
kubectl apply -f vllm-hpa.yaml -n ${NAMESPACE}
```
GKE は、デプロイする別の Pod をスケジュールします。これにより、ノードプールオートスケーラーがトリガーされ、2 番目の vLLM レプリカをデプロイする前に 2 番目のノードを追加します。

Pod の自動スケーリングの進行状況を確認します。

kubectl get hpa --watch -n ${NAMESPACE}