Trillium TPU에서 멀티 호스트 vLLM TorchTPU를 사용하여 GKE에서 Gemma 4 26B 서빙

이 튜토리얼에서는 멀티 호스트 Cloud TPU v6e (Trillium) 가속기와 vLLM TorchTPU를 사용하여 Google Kubernetes Engine (GKE)에서 멀티모달 전문가 혼합 모델 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에서 무료 계정에 가입할 수 있습니다. 계정이 있으면 액세스 토큰을 생성합니다.

  1. Hugging Face에 오신 것을 환영합니다 페이지에서 계정 아바타를 클릭하고 액세스 토큰을 선택합니다.
  2. 액세스 토큰 페이지에서 새 토큰 만들기를 클릭합니다.
  3. 읽기 토큰 유형을 선택하고 토큰 이름을 입력합니다.
  4. 액세스 토큰이 표시됩니다. 토큰을 안전한 장소에 저장합니다.
  • Google Cloud CLI를 설치하고 초기화하고 kubectl를 설치합니다. Cloud Shell을 사용하는 경우 두 도구 모두 사전 설치되어 있습니다.

튜토리얼을 완료하는 데 필요한 권한을 얻으려면 관리자에게 프로젝트에 대한 다음 IAM 역할을 부여해 달라고 요청하세요.

역할 부여에 대한 자세한 내용은 프로젝트, 폴더, 조직에 대한 액세스 관리를 참조하세요.

커스텀 역할이나 다른 사전 정의된 역할을 통해 필요한 권한을 얻을 수도 있습니다.

환경 변수 설정

환경 변수 설정:

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 노드 풀을 연결하려면 다음 단계를 완료합니다.

  1. 기본 노드 풀에 e2-standard-4 노드를 사용하고 2x4 토폴로지 (총 8개의 TPU 칩)로 정렬된 ct6e-standard-4t 노드 2개가 있는 TPU v6e 노드 풀을 사용하는 리전 GKE 클러스터를 만듭니다. 기본 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}"

    특정 예약을 사용하는 대신 다음 플래그를 전달하여 스팟 용량 또는 Flex-start VM을 사용하여 노드 풀을 만들 수도 있습니다.

    • 스팟 VM: --reservation-affinity 및 --reservation 플래그를 --spot로 바꿔 스팟 용량을 사용합니다. 스팟 용량을 사용하는 인스턴스는 노드 풀이 생성된 후 생성되는 데 시간이 걸릴 수 있으며 워크로드 실행 중에 선점될 수 있습니다. 포드를 만드는 데 시간이 걸리거나 실행 중에 워크로드가 중지되는 경우 이 문제가 원인일 수 있습니다.
    • Flex-start VM: --flex-start를 전달하여 Flex-start 소비 옵션을 사용하여 노드 풀을 프로비저닝합니다. 자세한 내용은 TPU 및 Flex-start VM으로 소규모 일괄 워크로드 실행을 참조하세요.
  2. 클러스터와 통신하도록 kubectl 사용자 인증 정보를 구성합니다.

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

사용자 인증 정보 구성 및 LeaderWorkerSet 설치

멀티 호스트 배포를 위해 클러스터를 준비하려면 다음 단계를 완료하세요.

  1. Hugging Face 토큰을 내보내고 hf-secret이라는 Kubernetes 보안 비밀을 만듭니다.

    kubectl create secret generic hf-secret \
        --from-literal=hf_api_token="${HF_TOKEN}" \
        --dry-run=client -o yaml | kubectl apply -f -
  2. 클러스터에 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)는 두 가지 리소스를 정의합니다.

  • 2x4 TPU 슬라이스의 두 ct6e-standard-4t 노드에 예약된 size: 2 (리더 포드 1개와 워커 포드 1개)가 있는 LeaderWorkerSet 리소스 (vllm-tpu-gemma4-26b) 각 포드는 --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를 실행합니다.
  • 리더 포드에서 포트 8000을 노출하는 ClusterIP 서비스 (vllm-service)

워크로드를 배포하고 모니터링하려면 다음 단계를 완료합니다.

  1. 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
  2. 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
  3. 애플리케이션 시작이 완료될 때까지 리더 포드 상태와 로그를 모니터링합니다.

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

    포드가 모델 가중치를 미리 가져오고, 8개의 TPU 칩 모두에서 2노드 Ray 클러스터를 형성하고, 모델을 컴파일하는 동안 초기 서비스 시작에는 약 9분이 걸립니다. 로그에 Application startup complete가 보고될 때까지 기다린 후 트래픽을 전송합니다.

  4. 두 번째 터미널 세션에서 로컬 포트 8000을 vllm-service에 전달합니다.

    kubectl port-forward service/vllm-service 8000:8000

    이 터미널 세션에서 포트 전달을 실행 중인 상태로 둡니다.

추론 요청 전송

첫 번째 터미널 세션으로 돌아가 텍스트 및 멀티모달 채팅 완료 요청을 모두 사용하여 배포된 모델을 테스트합니다.

  1. /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
        }'
  2. 이미지 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 클러스터를 삭제합니다.

  1. 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
  2. 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
  3. 두 번째 터미널 세션으로 돌아가 포트 전달을 중지합니다.

다음 단계