透過 vLLM 在 GKE 上使用 TPU Trillium 提供 LLM


本教學課程說明如何使用 vLLM 服務架構,在 Google Kubernetes Engine (GKE) 上,透過張量處理單元 (TPU) 提供大型語言模型 (LLM) 服務。在本教學課程中,您將提供 Llama 3.1 70b,使用 TPU Trillium,並透過 vLLM 伺服器指標設定水平 Pod 自動調度

如果您在部署及提供 AI/機器學習工作負載時,需要代管型 Kubernetes 的精細控制、擴充性、復原能力、可攜性和成本效益,這份文件就是不錯的入門資源。

背景

在 GKE 上使用 TPU Trillium,即可導入可用於正式環境的穩固服務解決方案,同時享有代管型 Kubernetes 的所有優點,包括高效擴充性和更高的可用性。本節說明本指南中使用的重要技術。

TPU Trillium

TPU 是 Google 開發的客製化特殊應用積體電路 (ASIC)。TPU 可加速處理使用 TensorFlowPyTorchJAX 等架構建構的機器學習和 AI 模型。本教學課程使用 Google 第六代 TPU「Trillium」。

在 GKE 中使用 TPU 之前,建議您先完成下列學習路徑:

  1. 瞭解 TPU Trillium 系統架構
  2. 瞭解 GKE 中的 TPU

vLLM

vLLM 是經過高度最佳化的開放原始碼框架,可提供 LLM。vLLM 可透過下列功能提高 TPU 的服務輸送量:

  • 使用 PagedAttention 實作最佳化轉換器。
  • 持續批次處理,提升整體服務輸送量。
  • 在多個 TPU 上進行張量平行處理和分散式服務。

詳情請參閱 vLLM 說明文件

Cloud Storage FUSE

Cloud Storage FUSE 可讓 GKE 叢集存取 Cloud Storage,以取得物件儲存空間值區中的模型權重。在本教學課程中,建立的 Cloud Storage bucket 一開始會是空的。vLLM 啟動時,GKE 會從 Hugging Face 下載模型,並將權重快取到 Cloud Storage bucket。Pod 重新啟動或部署作業擴大規模時,後續的模型載入作業會從 Cloud Storage 值區下載快取資料,並利用平行下載功能達到最佳效能。

詳情請參閱 Cloud Storage FUSE CSI 驅動程式說明文件

目標

本教學課程的適用對象為 MLOps 或 DevOps 工程師,以及平台管理員,他們想使用 GKE 自動化調度管理功能提供 LLM 服務。

本教學課程包含下列步驟:

  1. 根據模型特徵,建立具有建議 TPU Trillium 拓撲的 GKE 叢集。
  2. 在叢集的節點集區中部署 vLLM 架構。
  3. 使用負載平衡器,透過 vLLM 架構提供 Llama 3.1 70b。
  4. 使用 vLLM 伺服器指標設定水平 Pod 自動調度資源。
  5. 提供模型。

事前準備

  • 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

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Enable the API

  • Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin, roles/iam.securityAdmin, roles/artifactregistry.writer, roles/container.clusterAdmin

    Check for the roles

    1. In the Google Cloud console, go to the IAM page.

      Go to IAM
    2. Select the project.
    3. 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.

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

    1. In the Google Cloud console, go to the IAM page.

      前往 IAM
    2. 選取專案。
    3. 按一下「授予存取權」
    4. 在「New principals」(新增主體) 欄位中,輸入您的使用者 ID。 這通常是 Google 帳戶的電子郵件地址。

    5. 在「Select a role」(選取角色) 清單中,選取角色。
    6. 如要授予其他角色,請按一下 「新增其他角色」,然後新增每個其他角色。
    7. 按一下 [Save]
    8. 準備環境

      在本節中,您將佈建部署 vLLM 和模型所需的資源。

      取得模型存取權

      如要在 Hugging Face 存放區中使用 Llama 3.1 70b,請務必簽署同意聲明。

      產生存取權杖

      如果沒有,請產生新的 Hugging Face 權杖

      1. 依序點選「Your Profile」>「Settings」>「Access Tokens」
      2. 選取「New Token」
      3. 指定所選名稱和至少 Read 的角色。
      4. 選取「產生權杖」

      啟動 Cloud Shell

      在本教學課程中,您將使用 Cloud Shell 管理託管於Google Cloud的資源。Cloud Shell 已預先安裝本教學課程所需的軟體,包括 kubectlgcloud CLI

      如要使用 Cloud Shell 設定環境,請按照下列步驟操作:

      1. 在 Google Cloud 控制台中,按一下Cloud Shell 啟用圖示Google Cloud 控制台中的「啟用 Cloud Shell」,啟動 Cloud Shell 工作階段。系統會在 Google Cloud 控制台的底部窗格啟動工作階段。

      2. 設定預設環境變數:

        gcloud config set project PROJECT_ID && \
        gcloud config set billing/quota_project PROJECT_ID && \
        export PROJECT_ID=$(gcloud config get project) && \
        export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)") && \
        export CLUSTER_NAME=CLUSTER_NAME && \
        export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION && \
        export ZONE=ZONE && \
        export HF_TOKEN=HUGGING_FACE_TOKEN && \
        export CLUSTER_VERSION=CLUSTER_VERSION && \
        export GSBUCKET=GSBUCKET && \
        export KSA_NAME=KSA_NAME && \
        export NAMESPACE=NAMESPACE
        

        替換下列值:

        • PROJECT_ID:您的 Google Cloud 專案 ID
        • CLUSTER_NAME:GKE 叢集名稱。
        • CONTROL_PLANE_LOCATION:叢集控制層的 Compute Engine 區域。提供支援 TPU Trillium (v6e) 的地區。
        • ZONE:支援 TPU Trillium (v6e) 的可用區。
        • CLUSTER_VERSION:GKE 版本,必須支援您想使用的機器類型。請注意,預設 GKE 版本可能無法支援目標 TPU。GKE 1.31.2-gke.1115000 以上版本支援 TPU Trillium (v6e)。
        • GSBUCKET:用於 Cloud Storage FUSE 的 Cloud Storage 值區名稱。
        • KSA_NAME:用於存取 Cloud Storage bucket 的 Kubernetes ServiceAccount 名稱。Cloud Storage FUSE 必須具備 bucket 存取權才能運作。
        • NAMESPACE:您要部署 vLLM 資產的 Kubernetes 命名空間。

      建立 GKE 叢集

      您可以在 GKE Autopilot 或 Standard 叢集上,透過 TPU 供應 LLM。建議您使用 Autopilot 叢集,享受全代管 Kubernetes 體驗。如要為工作負載選擇最合適的 GKE 作業模式,請參閱「選擇 GKE 作業模式」。

      Autopilot

      1. 建立 GKE Autopilot 叢集:

        gcloud container clusters create-auto ${CLUSTER_NAME} \
            --cluster-version=${CLUSTER_VERSION} \
            --location=${CONTROL_PLANE_LOCATION}
        

      標準

      1. 建立 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
        
      2. 建立 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 建立下列資源:

      設定 kubectl 與叢集通訊

      如要設定 kubectl 與叢集通訊,請執行下列指令:

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

      為 Hugging Face 憑證建立 Kubernetes 密鑰

      1. 建立命名空間。如果使用 default 命名空間,可以略過這個步驟:

        kubectl create namespace ${NAMESPACE}
        
      2. 建立包含 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
      

      這會建立 Cloud Storage 值區,用於儲存從 Hugging Face 下載的模型檔案。

      設定 Kubernetes ServiceAccount 以存取 bucket

      1. 建立 Kubernetes ServiceAccount:

        kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}
        
      2. 授予 Kubernetes ServiceAccount 讀寫權限,以便存取 Cloud Storage bucket:

        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"
        
      3. 或者,您也可以授予專案中所有 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 bucket,用於儲存下載的模型和編譯快取。Cloud Storage FUSE CSI 驅動程式會讀取 bucket 的內容。
        2. 啟用檔案快取功能的磁碟區,以及 Cloud Storage FUSE 的平行下載功能。
        最佳做法

        視模型內容 (例如權重檔案) 的預期大小而定,使用 tmpfsHyperdisk / Persistent Disk 做為檔案快取的備份。在本教學課程中,您將使用以 RAM 為後端的 Cloud Storage FUSE 檔案快取。

      部署 vLLM 模型伺服器

      如要部署 vLLM 模型伺服器,本教學課程會使用 Kubernetes Deployment。Deployment 是 Kubernetes API 物件,可讓您執行多個 Pod 副本,並將這些副本分散到叢集的節點中。

      1. 檢查以下儲存為 vllm-llama3-70b.yaml 的 Deployment 資訊清單,該資訊清單使用單一副本:

        
        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: docker.io/vllm/vllm-tpu:73aa7041bfee43581314e6f34e9a657137ecc092
                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 錯誤。如要避免發生這項錯誤,有兩種做法:

        1. 從 Deployment YAML 中移除下列區塊,即可移除 XLA 快取位置。這表示所有副本都會重新編譯快取。

          - name: VLLM_XLA_CACHE_PATH
            value: "/data"
          
        2. 將 Deployment 擴展至 1,並等待第一個副本準備就緒並寫入 XLA 快取。然後擴充至其他副本。這樣一來,其餘副本就能讀取快取,而不必嘗試寫入快取。

      2. 執行下列指令來套用資訊清單:

        kubectl apply -f vllm-llama3-70b.yaml -n ${NAMESPACE}
        
      3. 查看執行中模型伺服器的記錄:

        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)
        

      提供模型

      1. 如要取得 VLLM 服務的外部 IP 位址,請執行下列指令:

        export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}' -n ${NAMESPACE})
        
      2. 使用 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 叢集預設會啟用這項功能。

      1. 在叢集上設定自訂指標 Stackdriver 轉接器:

        kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml
        
      2. 將「監控檢視者」角色新增至自訂指標 Stackdriver 轉接器使用的服務帳戶:

        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
        
      3. 將下列資訊清單儲存為 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
        
      4. 將其套用至叢集:

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

      在 vLLM 端點上建立負載

      在 vLLM 伺服器上建立負載,測試 GKE 如何使用自訂 vLLM 指標自動調度資源。

      1. 執行 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 替換為要執行的平行程序數量。

      2. 執行 bash 指令碼:

        chmod +x load.sh
        nohup ./load.sh &
        

      確認 Google Cloud Managed Service for Prometheus 是否會擷取指標

      Google Cloud Managed Service for Prometheus 擷取指標後,您就可以在 Cloud Monitoring 中查看指標,並將負載新增至 vLLM 端點。

      1. 前往 Google Cloud 控制台的「指標探索器」頁面。

        前往「Metrics Explorer」頁面

      2. 按一下「< > PromQL」

      3. 輸入下列查詢,觀察流量指標:

        vllm:num_requests_waiting{cluster='CLUSTER_NAME'}
        

      折線圖會顯示一段時間內測得的 vLLM 指標 (num_requests_waiting)。vLLM 指標會從 0 (預先載入) 向上調整至某個值 (載入後)。這張圖表會確認 vLLM 指標是否已擷取至 Google Cloud Managed Service for Prometheus。下圖顯示預先載入值從 0 開始,在一分鐘內達到接近 400 的最大載入後值。

      vllm:num_requests_waiting

      部署水平 Pod 自動調度資源設定

      決定要根據哪項指標自動調度資源時,建議您使用下列 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 副本擴充至兩個。如要使用 num_requests_waiting 部署水平 Pod 自動調度器設定,請按照下列步驟操作:

      1. 將下列資訊清單儲存為 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

      2. 部署水平 Pod 自動調度資源設定:

        kubectl apply -f vllm-hpa.yaml -n ${NAMESPACE}
        

        GKE 會排定部署另一個 Pod,這會觸發節點集區自動調度器新增第二個節點,然後部署第二個 vLLM 副本。

      3. 查看 Pod 自動調度資源的進度:

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

        輸出結果會與下列內容相似:

        NAME       REFERENCE             TARGETS       MINPODS   MAXPODS   REPLICAS   AGE
        vllm-hpa   Deployment/vllm-tpu   <unknown>/10   1         2         0          6s
        vllm-hpa   Deployment/vllm-tpu   34972m/10      1         2         1          16s
        vllm-hpa   Deployment/vllm-tpu   25112m/10      1         2         2          31s
        vllm-hpa   Deployment/vllm-tpu   35301m/10      1         2         2          46s
        vllm-hpa   Deployment/vllm-tpu   25098m/10      1         2         2          62s
        vllm-hpa   Deployment/vllm-tpu   35348m/10      1         2         2          77s
        
      4. 等待 10 分鐘,然後重複「確認 Google Cloud Managed Service for Prometheus 擷取指標」一節中的步驟。Google Cloud Managed Service for Prometheus 現在會從 vLLM 端點擷取指標。

      清除所用資源

      如要避免系統向您的 Google Cloud 帳戶收取本教學課程中所用資源的相關費用,請刪除含有該項資源的專案,或者保留專案但刪除個別資源。

      刪除已部署的資源

      如要避免系統向您的 Google Cloud 帳戶收取本指南所建立資源的費用,請執行下列指令:

      ps -ef | grep load.sh | awk '{print $2}' | xargs -n1 kill -9
      
      gcloud container clusters delete ${CLUSTER_NAME} \
        --location=${CONTROL_PLANE_LOCATION}
      

      後續步驟