Triển khai tính năng suy luận vLLM TPU nhiều máy chủ bằng Ray trên GKE

1. Giới thiệu

Trong lớp học lập trình này, bạn sẽ tìm hiểu cách triển khai các dịch vụ suy luận vLLM (Mô hình ngôn ngữ lớn ảo) nhiều máy chủ có hiệu suất cao trên Google Kubernetes Engine (GKE) bằng cách sử dụng TPU của Google Cloud. Bạn sẽ định cấu hình suy luận phân tán bằng Ray và quản lý khối lượng công việc một cách tự nhiên trên GKE bằng LeaderWorkerSets.

Hướng dẫn này mô phỏng một chế độ thiết lập sản xuất để phân phát các mô hình lớn như Qwen 30B.

Bạn sẽ thực hiện

  • Tạo một mạng VPC tuỳ chỉnh cho lưu lượng truy cập của trình tăng tốc.
  • Cung cấp một cụm GKE bằng Trình điều khiển Ray và trình điều khiển CSI GCS Fuse.
  • Khởi tạo một Bộ nhớ đệm nhanh GCS để tăng tốc độ tải mô hình.
  • Cung cấp một bộ nút TPU v6e có nhiều máy chủ với dung lượng được đặt trước.
  • Định cấu hình Workload Identity để truy cập an toàn vào trọng số mô hình.
  • Triển khai và kiểm thử công cụ vLLM phục vụ mô hình tham số 30B.

Bạn cần có

  • Một dự án trên Google Cloud đã bật tính năng thanh toán.
  • Google Cloud Reservation cho tài nguyên TPU phiên bản 6e (32 chip, ct6e-standard-4t).
  • Quyền truy cập để sao chép trọng số mô hình từ một nhóm nguồn.
  • Cloud Shell hoặc một cửa sổ dòng lệnh cục bộ đã cài đặt gcloud, kubectl và helm.
  • Thời lượng ước tính: 60 phút
  • Chi phí ước tính: Dưới 60 USD (giả sử việc tháo dỡ diễn ra ngay sau đó).

2. Trước khi bắt đầu

Tạo hoặc chọn một dự án trên Google Cloud

  1. Trong Google Cloud Console, hãy chọn hoặc tạo một dự án trên Google Cloud.
  2. Đảm bảo bạn đã bật tính năng thanh toán cho dự án trên đám mây.

Khởi động Cloud Shell

  1. Nhấp vào Kích hoạt Cloud Shell ở đầu bảng điều khiển Cloud.
  2. Xác minh hoạt động xác thực:
gcloud auth list
  1. Xác nhận dự án của bạn:
gcloud config get project
  1. Đặt nếu cần:
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

Đặt các biến môi trường

Để dễ dàng thực thi các lệnh, hãy xác định các biến sau trong shell. Thay thế <YOUR_ZONE> bằng vùng TPU được phân bổ và <YOUR_RESERVATION_NAME> bằng mã đặt trước. Bạn sẽ cần tạo Mã truy cập người dùng Hugging Face để tải các trọng số mô hình có kiểm soát quyền truy cập xuống. Sau khi tạo, hãy thay thế <YOUR_HUGGING_FACE_TOKEN> bằng mã thông báo bạn vừa tạo.

export PROJECT_ID=$(gcloud config get-value project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export ZONE="<YOUR_ZONE>" # e.g., us-east5-a
export REGION=${ZONE%-*}
export CLUSTER_NAME="qwen-serving-cluster"
export GVNIC_NETWORK_PREFIX="qwen-serving"
export BUCKET_NAME="inf-demo-model-storage-${PROJECT_NUMBER}"
export RESERVATION_NAME="<YOUR_RESERVATION_NAME>"
export NODE_POOL_NAME="tpu-v6e-32-resvd-pool"
export MULTIHOST_COLLECTION_NAME="tpu-6-collection"
export HF_TOKEN="<YOUR_HUGGING_FACE_TOKEN>" # Token with access to Qwen model if restricted

Bật API

Bật các dịch vụ bắt buộc của Google Cloud:

gcloud services enable \
    container.googleapis.com \
    compute.googleapis.com \
    iam.googleapis.com \
    cloudresourcemanager.googleapis.com

3. Tạo mạng tuỳ chỉnh

Các tải công việc TPU có nhiều máy chủ yêu cầu cấu hình mạng cụ thể, bao gồm cả kích thước MTU lớn hơn để giao tiếp hiệu quả với bộ tăng tốc. Tạo một mạng VPC tuỳ chỉnh cho cụm của bạn.

  1. Tạo mạng VPC có MTU lớn (8896):
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. Tạo mạng con cho cụm:
    gcloud compute --project=${PROJECT_ID} \
    networks subnets create ${GVNIC_NETWORK_PREFIX}-tpu \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --region=${REGION} \
    --range=192.168.100.0/24
  3. Tạo quy tắc tường lửa cho phép lưu lượng truy cập nội bộ để nhân viên có thể giao tiếp:
    gcloud compute --project=${PROJECT_ID} firewall-rules create ${GVNIC_NETWORK_PREFIX}-allow-internal \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --allow=all \
    --source-ranges=172.16.0.0/12,192.168.0.0/16,10.0.0.0/8 \
    --description="Allow all internal traffic within the network."

4. Cung cấp cụm GKE

Tạo một chế độ thiết lập cụm GKE tiêu chuẩn được định cấu hình để hỗ trợ các lượt gắn GCS Fuse và khối lượng công việc của Ray Operator.

  1. Tạo cụm:
    gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${REGION} \
    --release-channel=rapid \
    --machine-type=e2-standard-4 \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --subnetwork=${GVNIC_NETWORK_PREFIX}-tpu \
    --num-nodes=1 \
    --gateway-api=standard \
    --enable-managed-prometheus \
    --enable-dataplane-v2 \
    --enable-dataplane-v2-metrics \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --addons=GcsFuseCsiDriver,RayOperator \
    --enable-ip-alias
  2. Truy xuất thông tin xác thực của cụm:
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Tạo khoá bí mật Hugging Face: Lưu mã thông báo của bạn một cách an toàn để tải nội dung xuống từ vùng chứa:
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. Cài đặt LeaderWorkerSet (LWS) thông qua Helm. LWS quản lý các nhóm gồm những nhóm nhỏ phải được lên lịch cùng nhau:
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. Bật tính năng Bộ nhớ đệm nhanh của GCS

Để tăng tốc độ đọc hàng chục GB trọng số từ Cloud Storage trong quá trình phân phát, hãy tạo một bộ chứa GCS và bật GCS Rapid Cache trong vùng của bạn.

  1. Tạo vùng lưu trữ:
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. Khởi chạy Rapid Cache trong vùng TPU:
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. Thiết lập Workload Identity và quyền truy cập vào bộ nhớ

Định cấu hình mối liên kết danh tính để gắn an toàn nhóm trọng lượng vào các nhóm GKE mà không cần nhúng các khoá có thời gian tồn tại lâu dài.

  1. Tạo một Tài khoản dịch vụ IAM chuyên dụng:
    gcloud iam service-accounts create tpu-reader-sa
  2. Cấp quyền đọc cho bộ chứa:
    gcloud storage buckets add-iam-policy-binding gs://${BUCKET_NAME} \
    --member="serviceAccount:tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role="roles/storage.objectAdmin"
  3. Tạo Workload Identity Binding cho Tài khoản dịch vụ Kubernetes của không gian tên default:
    gcloud iam service-accounts add-iam-policy-binding tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com \
    --role="roles/iam.workloadIdentityUser" \
    --member="serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"
  4. Chú thích SA Kubernetes:
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. Thiết lập trọng số mô hình

Để phân phát mô hình có 30 tỷ tham số, bạn cần tải các trọng số xuống từ Hugging Face vào vùng lưu trữ GCS. Để bỏ qua hạn mức đĩa của Cloud Shell (5 GB), hãy sử dụng Job Kubernetes tiêu chuẩn để tải xuống trực tiếp trong cụm và ghi vào ổ đĩa GCS Fuse được gắn một cách an toàn.

  1. Triển khai Công việc tải mô hình xuống: Tạo và áp dụng tệp kê khai sau đây để bắt đầu tải xuống:
    cat <<EOF | kubectl apply -f -
apiVersion: batch/v1
kind: Job
metadata:
  name: model-downloader
spec:
  ttlSecondsAfterFinished: 60
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
        gke-gcsfuse/memory-limit: "0"
    spec:
      serviceAccountName: default
      restartPolicy: OnFailure
      containers:
      - name: downloader
        image: python:3.10-slim
        command: ["/bin/sh", "-c"]
        args:
        - |
          pip install -U "huggingface_hub[hf_transfer]" filelock
          export HF_HUB_ENABLE_HF_TRANSFER=1

          python -c '
          import filelock

          class DummyLock:
              def __init__(self, *args, **kwargs): pass
              def __enter__(self): return self
              def __exit__(self, *args): pass
              def acquire(self, *args, **kwargs): pass
              def release(self, *args, **kwargs): pass

          filelock.FileLock = DummyLock

          from huggingface_hub import snapshot_download
          snapshot_download(
              repo_id="Qwen/Qwen3-30B-A3B", 
              local_dir="/models/qwen3-30b-weights",
              local_dir_use_symlinks=False
          )
          '
        env:
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        volumeMounts:
        - name: model-weights
          mountPath: /models
      volumes:
      - name: model-weights
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: ${BUCKET_NAME}
            mountOptions: "implicit-dirs"
EOF
  2. Theo dõi quá trình tải xuống: Kiểm tra nhật ký của nhóm tải xuống để theo dõi tiến trình:
    kubectl logs -f job/model-downloader
    Chờ cho đến khi tác vụ hoàn tất với trạng thái thành công.

8. Tạo Nhóm nút TPU được đặt trước

Cung cấp lát TPU thực tế có nhiều máy chủ bằng cách sử dụng chế độ đặt trước dung lượng hiện có.

  1. Chạy lệnh tạo:
    gcloud beta container node-pools create ${NODE_POOL_NAME} \
    --project=${PROJECT_ID} \
    --cluster=${CLUSTER_NAME} \
    --region=${REGION} \
    --node-locations=${ZONE} \
    --machine-type=ct6e-standard-4t \
    --tpu-topology=4x8 \
    --num-nodes=8 \
    --scopes=https://www.googleapis.com/auth/cloud-platform \
    --reservation-affinity=specific \
    --reservation=${RESERVATION_NAME} \
    --accelerator-network-profile=auto \
    --node-labels=cloud.google.com/gke-nodepool-group-name=${MULTIHOST_COLLECTION_NAME} \
    --node-labels=cloud.google.com/gke-workload-type=HIGH_AVAILABILITY \
    --node-labels=cloud.google.com/gke-networking-dra-driver=true
  2. Chờ các nút tham gia: Bạn có thể quan sát quá trình mở rộng quy mô tổng hợp nút một cách trực tiếp. Chờ cho đến khi 8 nút chứa ct6e tham gia kubectl get nodes.

9. Triển khai dịch vụ vLLM

  1. Tạo Network Claims: Bạn cần yêu cầu môi trường mạng:
    cat <<EOF | kubectl apply -f -
apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: all-netdev
spec:
  spec:
    devices:
      requests:
      - name: req-netdev
        exactly:
          deviceClassName: netdev.google.com
          allocationMode: All
EOF
  2. Triển khai điểm cuối API của Bộ cân bằng tải:
    cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: Service
metadata:
  name: vllm-tpu-service
spec:
  type: LoadBalancer
  selector:
    leaderworkerset.sigs.k8s.io/name: vllm-tpu-qwen
    leaderworkerset.sigs.k8s.io/worker-index: "0"
  ports:
  - protocol: TCP
    port: 8000
    targetPort: 8000
EOF
  3. Triển khai khối lượng công việc LeaderWorkerSet: Tệp kê khai này bắt đầu tổng hợp đầu/worker Ray một cách linh động trên 8 máy chủ lát.
    cat <<EOF | kubectl apply -f -
apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: vllm-tpu-qwen
spec:
  replicas: 1
  leaderWorkerTemplate:
    size: 8
    restartPolicy: RecreateGroupOnPodRestart
    workerTemplate:
      metadata:
        annotations:
          gke-gcsfuse/volumes: "true"
          gke-gcsfuse/memory-limit: "0"
        labels:
          leaderworkerset.sigs.k8s.io/name: vllm-tpu-qwen
          gke-gcsfuse/volumes: "true"
      spec:
        hostname: vllm-tpu-qwen
        serviceAccountName: default
        containers:
        - name: vllm-tpu
          image: vllm/vllm-tpu:nightly
          command: ["sh", "-c"]
          args:
          - |
            MY_TPU_IP=\$(hostname -I | awk '{print \$1}')
            echo "My TPU Network IP is: \$MY_TPU_IP"

            LEADER_DNS="vllm-tpu-qwen-0.vllm-tpu-qwen"
            until getent hosts \$LEADER_DNS; do
              echo "DNS not ready. Sleeping 5s..."
              sleep 5
              done
            LEADER_IP=\$(getent hosts \$LEADER_DNS | awk '{print \$1}')

            export JAX_PLATFORMS=''
            export SCAN_TPU_CHIPS=True
            export TPU_MULTIHOST_BACKEND=ray
            export JAX_DISTRIBUTED_INITIALIZATION_TIMEOUT=300
            export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/lib
            export VLLM_HOST_IP=\$MY_TPU_IP

            if [ "\$LWS_WORKER_INDEX" = "0" ]; then
              echo "Starting Ray Head..."
              ray start --head --port=6379 --node-ip-address=\$MY_TPU_IP --resources='{"TPU": 4}' --block &
              sleep 20
              until ray status; do sleep 5; done

              echo "Starting vLLM API Server..."
              python3 -m vllm.entrypoints.openai.api_server \
                --model=/models/qwen3-30b-weights \
                --tensor-parallel-size=32 \
                --pipeline-parallel-size=1 \
                --distributed-executor-backend=ray \
                --host=0.0.0.0 --port=8000 \
                --enforce-eager \
                --gpu-memory-utilization=0.90
            else
              ray start --address=\${LEADER_IP}:6379 --node-ip-address=\$MY_TPU_IP --resources='{"TPU": 4}' --block
            fi
          ports:
          - containerPort: 8000
          - containerPort: 6379
          volumeMounts:
          - name: model-weights
            mountPath: /models
            readOnly: true
          - name: dshm
            mountPath: /dev/shm
          resources:
            claims:
            - name: net-resources
            limits:
              google.com/tpu: 4
              memory: "100Gi"
            requests:
              google.com/tpu: 4
              memory: "100Gi"
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x8
          gke.networks.io/accelerator-network-profile: auto
        resourceClaims:
        - name: net-resources
          resourceClaimTemplateName: all-netdev
        volumes:
        - name: model-weights
          csi:
            driver: gcsfuse.csi.storage.gke.io
            volumeAttributes:
              bucketName: ${BUCKET_NAME}
              mountOptions: "implicit-dirs"
        - name: dshm
          emptyDir:
            medium: Memory
EOF

10. Phản hồi về việc triển khai thử nghiệm

Có thể mất từ 5 đến 10 phút để tất cả các nhóm trong LeaderWorkerSet kéo hình ảnh vùng chứa, khởi động Ray và trở thành Ready hoàn toàn. Bạn có thể theo dõi trạng thái bằng cách xem quá trình khởi động nhóm:

kubectl get pods -l leaderworkerset.sigs.k8s.io/name=vllm-tpu-qwen -w

Đợi cho đến khi tất cả 8 nhóm vllm-tpu-qwen- đều hiển thị STATUS dưới dạng Running và READY dưới dạng 2/2, đồng thời đảm bảo trình cân bằng tải đã nhận được một IP ngoài trước khi tiếp tục. Quá trình này có thể mất từ 7 đến 10 phút.

  1. Truy xuất IP ngoài: 
    export EXTERNAL_IP=$(kubectl get svc vllm-tpu-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
echo $EXTERNAL_IP

Thận trọng: Trong một dịch vụ sản xuất, điểm cuối này phải được bảo mật bằng một thứ gì đó như Identity Aware Proxy (IAP)

  1. Gửi yêu cầu suy luận bằng curl:
        curl -N -s http://$EXTERNAL_IP:8000/v1/chat/completions \
        -H "Content-Type: application/json" \
        -d '{
            "model": "/models/qwen3-30b-weights",
            "messages": [{"role": "user", "content": "Write a haiku about high-performance computing on TPUs."}],
            "temperature": 0.7,
            "max_tokens": 100,
            "stream": true
        }' | sed 's/^data: //' | grep -v '\[DONE\]' | grep -v '^$' | jq -rj '.choices[0].delta.content // empty' ; echo ""
    Bạn sẽ thấy đầu ra tương tự như một phản hồi JSON chứa văn bản thể hiện suy luận đã tạo của bạn!

11. Dọn dẹp

Để tránh các khoản phí phát sinh cho tài khoản Google Cloud của bạn, hãy xoá các tài nguyên đã tạo trong lớp học lập trình này.

  1. Xoá nhóm nút:
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. Xoá nhóm:
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. Xoá chế độ thiết lập Mạng và tường lửa:
    gcloud compute firewall-rules delete \
    "${GVNIC_NETWORK_PREFIX}-allow-internal" \
    --project="${PROJECT_ID}" --quiet

gcloud compute networks subnets delete "${GVNIC_NETWORK_PREFIX}-tpu" \
    --region="${REGION}" --quiet

gcloud compute networks delete "${GVNIC_NETWORK_PREFIX}-main" --quiet
  4. Huỷ liên kết và xoá tài khoản dịch vụ:
        # 1. Create the cleanup script
    cat << 'EOF' > clean_up_sa.sh
    #!/bin/bash

    # Validate that PROJECT_ID is available
    if [ -z "$PROJECT_ID" ]; then
      echo "Error: PROJECT_ID environment variable is not set."
      exit 1
    fi

    SA_EMAIL="tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com"
    SA_MEMBER="serviceAccount:${SA_EMAIL}"

    echo "Gathering IAM policy for ${SA_EMAIL}..."

    # Fetch roles assigned to this specific SA
    ROLES=$(gcloud projects get-iam-policy ${PROJECT_ID} \
        --flatten="bindings[].members" \
        --filter="bindings.members:${SA_MEMBER}" \
        --format="value(bindings.role)")

    if [ -z "$ROLES" ]; then
        echo "No IAM bindings found for this service account."
    else
        for ROLE in $ROLES; do
            echo "Removing binding for: ${ROLE}..."
            gcloud projects remove-iam-policy-binding ${PROJECT_ID} \
                --member="${SA_MEMBER}" \
                --role="${ROLE}" --quiet > /dev/null
        done
        echo "Successfully unbound all roles."
    fi

    # 2. Delete the service account itself
    echo "Deleting service account..."
    gcloud iam service-accounts delete ${SA_EMAIL} --project=${PROJECT_ID} --quiet

    echo "Cleanup complete."
    EOF

    # 2. Make the script executable and run it
    chmod +x clean_up_sa.sh
    ./clean_up_sa.sh
  5. Xoá bộ chứa GCS Chuyển đến bảng điều khiển đám mây, chọn Cloud Storage -> Buckets (Bộ chứa), chọn inf-demo-model-storage rồi chọn "Xoá".

12. Xin chúc mừng

Xin chúc mừng! Bạn đã triển khai thành công một ngăn xếp vLLM có tốc độ suy luận cao gồm nhiều máy chủ, sử dụng Ray một cách tự nhiên trên Google Kubernetes Engine.

Kiến thức bạn học được

  • Cung cấp các đường dẫn tuỳ chỉnh phù hợp với lưu lượng truy cập TPU tốc độ cao.
  • Tăng trọng số bằng cách sử dụng GCS Fuse và bộ nhớ đệm nhanh theo khu vực.
  • Điều phối các lát khối lượng công việc có nhiều máy chủ được đồng bộ hoá một cách tự nhiên thông qua LeaderWorkerSets.
  • Để tìm hiểu thêm, hãy xem Hướng dẫn sử dụng vLLM và Hướng dẫn triển khai llm-d