Google は AI 技術を使用して、コンテンツをご希望の言語に翻訳しています。AI 翻訳には誤りが含まれる場合があります。

Trillium TPU でマルチホスト vLLM TorchTPU を使用して GKE で Gemma 4 26B をサービングする

このチュートリアルでは、マルチホスト Cloud TPU v6e（Trillium）アクセラレータと vLLM TorchTPU を使用して、Google Kubernetes Engine（GKE）にマルチモーダル Mixture of Experts モデル Gemma 4 26B（google/gemma-4-26B-A4B-it）をデプロイして提供する方法について説明します。ワークロードは 2x4 トポロジの複数の TPU ホストにまたがるため、Kubernetes LeaderWorkerSet API と Ray を使用して、8 つの TPU チップに分散されたテンソル並列推論を調整します。

目標

GKE クラスタとマルチホスト TPU v6e ノードプールを設定します。
LeaderWorkerSet を使用してマルチホスト vLLM サーバーをデプロイします。
テキストとマルチモーダルの推論リクエストを送信する。
サービングベンチマークワークロードを実行します。
クリーンアップする。

費用

このドキュメントでは、課金対象である次のGoogle Cloudコンポーネントを使用します。

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。

このドキュメントに記載されているタスクの完了後、作成したリソースを削除すると、それ以上の請求は発生しません。詳細については、クリーンアップをご覧ください。

始める前に

このチュートリアルを使用するには、Hugging Face アクセストークンが必要です。Hugging Face で無料アカウントに登録できます。アカウントを作成したら、アクセストークンを生成します。

[Welcome to Hugging Face] ページで、アカウントのアバターをクリックし、[アクセストークン] を選択します。
[アクセストークン] ページで、[新しいトークンを作成] をクリックします。
[Read] トークンタイプを選択し、トークンの名前を入力します。
アクセストークンが表示されます。トークンを安全な場所に保存します。

Google Cloud CLI をインストールして初期化し、kubectl をインストールします。Cloud Shell を使用する場合、どちらのツールもプリインストールされています。

このチュートリアルを完了するために必要な権限を取得するには、プロジェクトに対する次の IAM ロールを付与するよう管理者に依頼してください。

Kubernetes Engine 管理者（roles/container.admin）
Compute 管理者（roles/compute.admin）
サービスアカウントユーザー（roles/iam.serviceAccountUser）

ロールの付与については、プロジェクト、フォルダ、組織へのアクセス権の管理をご覧ください。

必要な権限は、カスタムロールや他の事前定義ロールから取得することもできます。

環境変数を設定する

環境変数を設定します。

export PROJECT_ID="YOUR_PROJECT_ID"
export RESERVATION_NAME="YOUR_RESERVATION_NAME"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export NODE_POOL_NAME="tpuv6e-pool"
export VLLM_IMAGE="vllm/vllm-torchtpu:v0.29.0"

gcloud config set project "${PROJECT_ID}"
gcloud config set billing/quota_project "${PROJECT_ID}"

次のように置き換えます。

YOUR_PROJECT_ID: 実際の Google Cloudプロジェクトの ID。
YOUR_RESERVATION_NAME: 予約の名前。
YOUR_REGION: GKE クラスタをデプロイするリージョン。
YOUR_ZONE: TPU v6e ノードプールをデプロイするゾーン。
YOUR_CLUSTER_NAME: GKE クラスタの名前。
YOUR_HF_TOKEN: Hugging Face アクセストークン。
YOUR_VLLM_TPU_IMAGE: 使用する vLLM TPU イメージ。公開イメージについては、vLLM TorchTPU リポジトリをご覧ください。

GKE クラスタを作成する

GKE クラスタを作成してマルチホスト TPU v6e ノードプールを接続する手順は次のとおりです。

デフォルトのノードプールに e2-standard-4 ノードを使用するリージョン GKE クラスタと、2x4 トポロジに配置された 2 つの ct6e-standard-4t ノード（合計 8 つの TPU チップ）を持つ TPU v6e ノードプールを作成します。デフォルトの GKE マシンタイプ（e2-medium）では、LeaderWorkerSet コントローラに十分な CPU リソースが提供されないため、このチュートリアルでは e2-standard-4 を指定します。
```
gcloud container clusters create "${CLUSTER_NAME}" \
    --project="${PROJECT_ID}" \
    --location="${REGION}" \
    --machine-type=e2-standard-4 \

gcloud container node-pools create "${NODE_POOL_NAME}" \
    --project="${PROJECT_ID}" \
    --location="${REGION}" \
    --cluster="${CLUSTER_NAME}" \
    --node-locations="${ZONE}" \
    --machine-type=ct6e-standard-4t \
    --tpu-topology=2x4 \
    --num-nodes=2 \
    --disk-size=500 \
    --reservation-affinity=specific \
    --reservation="${RESERVATION_NAME}"
```
特定の予約を使用する代わりに、次のフラグを渡して、Spot 容量または Flex Start VM を使用してノードプールを作成することもできます。
- Spot VM: --reservation-affinity フラグと --reservation フラグを --spot に置き換えて、Spot 容量を使用します。Spot 容量を使用するインスタンスは、ノードプールの作成後に作成されるまでに時間がかかることがあり、ワークロードの実行中にプリエンプトされる可能性があります。Pod の作成に時間がかかる場合や、ワークロードが実行中に停止する場合は、これが原因である可能性があります。
- Flex Start VM: --flex-start を渡して、Flex Start 消費オプションを使用してノードプールをプロビジョニングします。詳細については、TPU と Flex Start VM で小規模なバッチワークロードを実行するをご覧ください。

クラスタと通信するように kubectl 認証情報を構成します。

gcloud container clusters get-credentials "${CLUSTER_NAME}" \
    --project="${PROJECT_ID}" \
    --location="${REGION}"

認証情報を構成して LeaderWorkerSet をインストールする

マルチホストデプロイ用にクラスタを準備するには、次の手順を完了します。

Hugging Face トークンをエクスポートし、hf-secret という名前の Kubernetes Secret を作成します。

kubectl create secret generic hf-secret \
    --from-literal=hf_api_token="${HF_TOKEN}" \
    --dry-run=client -o yaml | kubectl apply -f -

クラスタに LeaderWorkerSet CustomResourceDefinition とコントローラをインストールします。

kubectl apply --server-side -k "github.com/kubernetes-sigs/lws/config/default?ref=main"
echo "Waiting for LWS controller to be ready..."
kubectl wait --for=condition=Available --timeout=300s deployment/lws-controller-manager -n lws-system

マルチホスト vLLM サーバーをデプロイする

ワークロードマニフェスト（gemma4-26b-multihost-v6e-2x4.yaml）は、次の 2 つのリソースを定義します。

2x4 TPU スライス内の 2 つの ct6e-standard-4t ノードにスケジュールされた size: 2（1 つのリーダー Pod と 1 つのワーカー Pod）を含む LeaderWorkerSet リソース（vllm-tpu-gemma4-26b）。各 Pod は、--tensor-parallel-size=8、--kv-cache-dtype=fp8、--max-model-len=5248、--max-num-seqs=128 フラグを使用して Ray と vLLM OpenAI API サーバーを起動する前に、Hugging Face から google/gemma-4-26B-A4B-it モデルの重みをプリフェッチしてインメモリボリュームに格納する InitContainer を実行します。
リーダー Pod のポート 8000 を公開する ClusterIP Service（vllm-service）。

ワークロードをデプロイしてモニタリングするには、次の手順を完了します。

GitHub から LeaderWorkerSet マニフェストを取得します。

curl -LO https://raw.githubusercontent.com/GoogleCloudSamples/AIHypercomputerSamples/main/tpu/inference/gemma-4-26b/gemma4-26b-multihost-v6e-2x4.yaml

または、次の内容の gemma4-26b-multihost-v6e-2x4.yaml という名前のファイルを作成します。

apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: vllm-tpu-gemma4-26b
spec:
  replicas: 1
  leaderWorkerTemplate:
    # One pod per TPU host: a 2x4 v6e slice has 2 ct6e-standard-4t hosts.
    size: 2
    restartPolicy: RecreateGroupOnPodRestart
    leaderTemplate:
      metadata:
        labels:
          role: leader
      spec:
        # Download the model into the pod-local cache before vLLM starts, so the
        # serving ranks can run offline and never hit Hugging Face during init.
        initContainers:
        - name: prefetch-model
          image: $VLLM_IMAGE
          command: ["hf", "download", "google/gemma-4-26B-A4B-it"]
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: HF_HOME
            value: "/data/huggingface_cache"
          volumeMounts:
          - name: model-cache
            mountPath: /data
        containers:
        - name: vllm-leader
          image: $VLLM_IMAGE
          command:
            - sh
            - -c
            - |
              set -e
              mkdir -p "${TMPDIR}"
              ray start --head --port=6379
              echo "Waiting for all 8 TPU chips to join the Ray cluster..."
              until ray status 2>/dev/null | grep -E -q "8(\.0)? TPU"; do
                sleep 5
              done
              vllm serve google/gemma-4-26B-A4B-it \
                --tensor-parallel-size 8 \
                --distributed-executor-backend ray \
                --host 0.0.0.0 \
                --port 8000 \
                --max-model-len 5248 \
                --max-num-seqs 128 \
                --max-num-batched-tokens 4096 \
                --kv-cache-dtype fp8 \
                --gpu-memory-utilization 0.85 \
                --disable-chunked-mm-input
          resources:
            requests:
              cpu: "128"
              memory: "384Gi"
              google.com/tpu: 4
            limits:
              cpu: "170"
              memory: "650Gi"
              google.com/tpu: 4
          ports:
          - containerPort: 8000
          - containerPort: 6379
          env:
          - name: TPU_MULTIHOST_BACKEND
            value: "ray"
          - name: VLLM_ENGINE_READY_TIMEOUT_S
            value: "3600"
          # Same cache the init container populated; offline mode guarantees no
          # Hugging Face requests from the serving ranks.
          - name: HF_HOME
            value: "/data/huggingface_cache"
          - name: HF_HUB_OFFLINE
            value: "1"
          - name: TMPDIR
            value: "/data/tmp"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
          - mountPath: /data
            name: model-cache
          # Model loading and compilation can take a long time; allow up to 60 minutes.
          startupProbe:
            httpGet:
              path: /health
              port: 8000
            periodSeconds: 10
            failureThreshold: 360
          livenessProbe:
            httpGet:
              path: /health
              port: 8000
            periodSeconds: 10
          readinessProbe:
            httpGet:
              path: /health
              port: 8000
            periodSeconds: 5
        nodeSelector:
          topology.kubernetes.io/zone: "$ZONE"
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 2x4
          cloud.google.com/reservation-name: $RESERVATION_NAME
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 32Gi
        - name: model-cache
          emptyDir:
            medium: Memory
            sizeLimit: 350Gi
    workerTemplate:
      spec:
        # Each pod has its own cache, so the worker downloads the model too.
        initContainers:
        - name: prefetch-model
          image: $VLLM_IMAGE
          command: ["hf", "download", "google/gemma-4-26B-A4B-it"]
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: HF_HOME
            value: "/data/huggingface_cache"
          volumeMounts:
          - name: model-cache
            mountPath: /data
        containers:
        - name: vllm-worker
          image: $VLLM_IMAGE
          command:
            - sh
            - -c
            - |
              set -e
              mkdir -p "${TMPDIR}"
              echo "Waiting for the Ray head at ${LWS_LEADER_ADDRESS}:6379..."
              until ray status --address="${LWS_LEADER_ADDRESS}:6379" 2>/dev/null | grep -q "node_"; do
                sleep 3
              done
              ray start --address="${LWS_LEADER_ADDRESS}:6379" --block
          resources:
            requests:
              cpu: "128"
              memory: "384Gi"
              google.com/tpu: 4
            limits:
              cpu: "170"
              memory: "650Gi"
              google.com/tpu: 4
          env:
          - name: TPU_MULTIHOST_BACKEND
            value: "ray"
          - name: VLLM_ENGINE_READY_TIMEOUT_S
            value: "3600"
          - name: HF_HOME
            value: "/data/huggingface_cache"
          - name: HF_HUB_OFFLINE
            value: "1"
          - name: TMPDIR
            value: "/data/tmp"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
          - mountPath: /data
            name: model-cache
        nodeSelector:
          topology.kubernetes.io/zone: "$ZONE"
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 2x4
          cloud.google.com/reservation-name: $RESERVATION_NAME
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 32Gi
        - name: model-cache
          emptyDir:
            medium: Memory
            sizeLimit: 350Gi
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-service
spec:
  selector:
    leaderworkerset.sigs.k8s.io/name: vllm-tpu-gemma4-26b
    role: leader
  type: ClusterIP
  ports:
    - protocol: TCP
      port: 8000
      targetPort: 8000

LeaderWorkerSet マニフェストを適用します。

MANIFEST="gemma4-26b-multihost-v6e-2x4.yaml"
echo "Applying LeaderWorkerSet (retrying if webhook is not fully ready)..."
MAX_ATTEMPTS=12
attempt=1
until envsubst '$RESERVATION_NAME $ZONE $VLLM_IMAGE' < "${MANIFEST}" | kubectl apply -f -; do
  if [ "${attempt}" -ge "${MAX_ATTEMPTS}" ]; then
    echo "Error: failed to apply ${MANIFEST} after ${MAX_ATTEMPTS} attempts." >&2
    exit 1
  fi
  echo "Webhook not ready yet, retrying in 5 seconds (attempt ${attempt}/${MAX_ATTEMPTS})..."
  attempt=$((attempt + 1))
  sleep 5
done

アプリケーションの起動が完了するまで、リーダー Pod のステータスとログをモニタリングします。

LIMIT=60
count=0
until kubectl get pod -l role=leader -o name | grep -q "vllm-tpu"; do
  if [ "${count}" -ge "${LIMIT}" ]; then
    echo "Error: timeout waiting for leader pod to be created." >&2
    exit 1
  fi
  echo "Waiting for leader pod to be created..."
  sleep 5
  count=$((count + 1))
done

# Wait up to 20 minutes for the pod to start running (TPU nodes take a while to boot)
echo "Waiting for pod to be scheduled and running..."
if ! kubectl wait --for=jsonpath='{.status.phase}'=Running --timeout=1200s pod -l role=leader; then
  echo "Pod failed to start running within 20 minutes! Here are the events:" >&2
  kubectl describe pod -l role=leader
  exit 1
fi

echo "Pod is running. Streaming logs while waiting for it to become Ready (this includes model download)..."
# Stream logs in the background so the user can see what's happening
kubectl logs -f -l role=leader &
LOG_PID=$!
trap 'kill "${LOG_PID}" 2>/dev/null || true' EXIT

if ! kubectl wait --for=condition=Ready --timeout=3600s pod -l role=leader; then
  echo "Pod failed to become Ready within 60 minutes!" >&2
  kubectl describe pod -l role=leader
  exit 1
fi

echo "Pod is Ready!"

Pod がモデルの重みをプリフェッチし、8 つの TPU チップすべてに 2 ノードの Ray クラスタを形成して、モデルをコンパイルする間、最初のサービス起動には約 9 分かかります。トラフィックを送信する前に、ログに Application startup complete が報告されるまで待ちます。

2 番目のターミナルセッションで、ローカルポート 8000 を vllm-service に転送します。
```
kubectl port-forward service/vllm-service 8000:8000
```
このターミナルセッションでポート転送を実行したままにします。

推論リクエストを送信する

最初のターミナルセッションに戻り、テキストとマルチモーダルチャットの完了リクエストの両方を使用して、デプロイされたモデルをテストします。

/v1/chat/completions エンドポイントにテキストチャット完了リクエストを送信します。

curl http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "google/gemma-4-26B-A4B-it",
        "messages": [{"role": "user", "content": "Explain Mixture of Experts in 2 sentences."}],
        "max_tokens": 128
    }'

画像 URL とテキストプロンプトの両方を含むマルチモーダルリクエストを送信します。

curl http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "google/gemma-4-26B-A4B-it",
        "messages": [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": "https://images.unsplash.com/photo-1501594907352-04cda38ebc29?auto=format&fit=crop&w=800&q=80"
                        }
                    },
                    {
                        "type": "text",
                        "text": "Describe the image above"
                    }
                ]
            }
        ],
        "max_tokens": 300,
        "temperature": 0.0,
        "top_p": 1.0
    }'

ベンチマークワークロードを実行する

デプロイされたモデルをベンチマークするには、次のコマンドを実行します。

kubectl exec vllm-tpu-gemma4-26b-0 -- vllm bench serve \
  --backend openai-chat \
  --endpoint /v1/chat/completions \
  --model google/gemma-4-26B-A4B-it \
  --host localhost \
  --port 8000 \
  --dataset-name random \
  --num-prompts 250 \
  --random-input-len 1024 \
  --random-output-len 1024 \
  --max-concurrency 128

クリーンアップ

このチュートリアルで使用したリソースについて Google Cloud アカウントに継続的な料金が発生しないようにするには、Kubernetes ワークロード、TPU ノードプール、GKE クラスタを削除します。

vLLM サーバーのデプロイと Hugging Face シークレットを削除します。

envsubst '$RESERVATION_NAME $ZONE $VLLM_IMAGE' < gemma4-26b-multihost-v6e-2x4.yaml | kubectl delete -f - --ignore-not-found || true
kubectl delete secret hf-secret --ignore-not-found || true

TPU v6e ノードプールと GKE クラスタを削除します。

gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --location="${REGION}" \
    --project="${PROJECT_ID}" \
    --quiet || true

gcloud container clusters delete "${CLUSTER_NAME}" \
    --location="${REGION}" \
    --project="${PROJECT_ID}" \
    --quiet

2 つ目のターミナルセッションに戻り、ポート転送を停止します。

次のステップ

vLLM TorchTPU の詳細を確認する。
Cloud TPU の詳細を確認する。