使用 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 用于加速使用 TensorFlow、PyTorch 和 JAX 等框架构建的机器学习和 AI 模型。本教程使用 TPU Trillium，这是 Google 的第六代 TPU。

使用 GKE 中的 TPU 之前，我们建议您完成以下学习路线：

  1. 了解 TPU Trillium 系统架构。
  2. 了解 GKE 中的 TPU。

vLLM

vLLM 是一个经过高度优化的开源框架，用于提供 LLM。vLLM 可提高 GPU 上的服务吞吐量，具有如下功能：

  • 具有 PagedAttention 且经过优化的 Transformer（转换器）实现。
  • 连续批处理，可提高整体服务吞吐量。
  • 多个 TPU 上的张量并行处理和分布式服务。

如需了解详情，请参阅 vLLM 文档。

Cloud Storage FUSE

Cloud Storage FUSE 可让您从 GKE 集群访问位于对象存储分区中的模型权重。在本教程中，创建的 Cloud Storage 存储桶最初将是空的。当 vLLM 启动时，GKE 会从 Hugging Face 下载模型，并将权重缓存到 Cloud Storage 存储桶。在 Pod 重启或部署扩容时，后续的模型加载将从 Cloud Storage 存储桶下载缓存的数据，利用并行下载实现最佳性能。

如需了解详情，请参阅 Cloud Storage FUSE CSI 驱动程序文档。

目标

本教程适用于希望使用 GKE 编排功能提供 LLM 的 MLOps 或 DevOps 工程师或是平台管理员。

本教程介绍以下步骤：

  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

  • Make sure 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

  • Make sure 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 colunn 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. 在新的主账号字段中，输入您的用户标识符。 这通常是 Google 账号的电子邮件地址。

    5. 在选择角色列表中，选择一个角色。
    6. 如需授予其他角色，请点击 添加其他角色，然后添加其他各个角色。
    7. 点击 Save（保存）。

准备环境

在本部分中，您将预配部署 vLLM 和模型所需的资源。

获取对模型的访问权限

您必须签署同意协议，才能使用 Hugging Face 代码库中的 Llama 3.1 70b。

生成一个访问令牌

如果您还没有 Hugging Face 令牌，请生成一个新令牌：

  1. 点击您的个人资料 > 设置 > 访问令牌。
  2. 选择新建令牌 (New Token)。
  3. 指定您选择的名称和一个至少为 Read 的角色。
  4. 选择生成令牌。

启动 Cloud Shell

在本教程中，您将使用 Cloud Shell 来管理Google Cloud上托管的资源。Cloud Shell 预安装了本教程所需的软件，包括 kubectl 和 gcloud CLI。

如需使用 Cloud Shell 设置您的环境，请按照以下步骤操作：

  1. 在 Google Cloud 控制台中，点击 Google Cloud 控制台中的 Cloud Shell 激活图标 激活 Cloud Shell 以启动 Cloud Shell 会话。此操作会在 Google Cloud 控制台的底部窗格中启动会话。

  2. 设置默认环境变量：

    gcloud config set 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 ZONE=ZONE && \
export HF_TOKEN=HUGGING_FACE_TOKEN && \
export CLUSTER_VERSION=CLUSTER_VERSION && \
export GSBUCKET=GSBUCKET && \
export KSA_NAME=KSA_NAME && \
export NAMESPACE=NAMESPACE \
export IMAGE_NAME=IMAGE_NAME

    替换以下值：

    • PROJECT_ID：您的 Google Cloud 项目 ID。
    • CLUSTER_NAME：GKE 集群的名称。
    • ZONE：支持 TPU Trillium (v6e) 的可用区。
    • CLUSTER_VERSION：GKE 版本，必须支持您要使用的机器类型。请注意，默认 GKE 版本可能无法为您的目标 TPU 提供可用性。GKE 1.31.2-gke.1115000 或更高版本支持 TPU Trillium。
    • GSBUCKET：要用于 Cloud Storage FUSE 的 Cloud Storage 存储桶的名称。
    • KSA_NAME：用于访问 Cloud Storage 存储分区的 Kubernetes ServiceAccount 的名称。需要存储分区访问权限才能使用 Cloud Storage FUSE。
    • NAMESPACE：您要部署 vLLM 资产的 Kubernetes 命名空间。
    • IMAGE_NAME：vLLM TPU 映像。您可以使用公共映像 docker.io/vllm/vllm-tpu:2e33fe419186c65a18da6668972d61d7bbc31564，也可以自行构建映像。

创建 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 \
    --release-channel=rapid

    将 PROJECT_ID 替换为您的 Google Cloud 项目 ID。

标准

  1. 创建 GKE Standard 集群：

    gcloud container clusters create CLUSTER_NAME \
    --project=PROJECT_ID \
    --zone=ZONE \
    --cluster-version=CLUSTER_VERSION \
    --release-channel=rapid \
    --workload-pool=PROJECT_ID.svc.id.goog \
    --addons GcsFuseCsiDriver

  2. 创建 TPU 分片节点池：

    gcloud container node-pools create tpunodepool \
    --zone=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=ZONE

为 Hugging Face 凭据创建 Kubernetes Secret

  1. 创建一个命名空间。如果您使用的是 default 命名空间，则可以跳过此步骤：

    kubectl create namespace NAMESPACE

  2. 创建包含 Hugging Face 令牌的 Kubernetes Secret，运行以下命令：

    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=HUGGING_FACE_TOKEN \
    --namespace NAMESPACE

创建 Cloud Storage 存储桶

在 Cloud Shell 中，运行以下命令：

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

这会创建一个 Cloud Storage 存储桶来存储您从 Hugging Face 下载的模型文件。

设置 Kubernetes ServiceAccount 以访问存储桶

  1. 创建 Kubernetes ServiceAccount：

    kubectl create serviceaccount KSA_NAME --namespace NAMESPACE

  2. 向 Kubernetes ServiceAccount 授予读写权限，以便访问 Cloud Storage 存储桶：

    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 存储桶。Cloud Storage FUSE CSI 驱动程序会读取存储桶中的内容。
    2. 启用了文件缓存的卷以及 Cloud Storage FUSE 的并行下载功能。
    最佳实践：

    根据模型内容（例如权重文件）的预期大小，使用由 tmpfs 或 Hyperdisk / Persistent Disk 支持的文件缓存。在本教程中，您将使用由 RAM 支持的 Cloud Storage FUSE 文件缓存。

（可选）构建并部署 TPU 映像

如果您需要对 Docker 映像的内容进行更精细的控制，请选择此选项。

将 vLLM 服务器容器化：

  1. 克隆 vLLM 代码库并构建映像：

    git clone https://github.com/vllm-project/vllm && cd vllm && git reset --hard 2e33fe419186c65a18da6668972d61d7bbc31564 && docker build -f Dockerfile.tpu . -t vllm-tpu

  2. 将该映像推送到 Artifact Registry。

    gcloud artifacts repositories create vllm-tpu --repository-format=docker --location=REGION_NAME && \
gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
docker image tag vllm-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-tpu/vllm-tpu:latest && \
docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/vllm-tpu/vllm-tpu:latest

部署 vLLM 模型服务器

如需部署 vLLM 模型服务器，请按以下步骤操作：

  1. 检查保存为 vllm-llama3-70b.yaml 的部署清单。Deployment 是一种 Kubernetes API 对象，可让您运行分布在集群节点中的 Pod 的多个副本：

    
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: IMAGE_NAME
        command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
        args:
        - --host=0.0.0.0
        - --port=8000
        - --tensor-parallel-size=8
        - --max-model-len=8192
        - --model=meta-llama/Meta-Llama-3.1-70B
        - --download-dir=/data
        env: 
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        - name: VLLM_XLA_CACHE_PATH
          value: "/data"
        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

  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. 运行以下命令可获取服务的外部 IP 地址：

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

  2. 在新的终端中，使用 curl 与模型互动：

    curl http://$vllm_service:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
    "model": "meta-llama/Meta-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/Meta-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. 将 Monitoring Viewer 角色添加到自定义指标 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

在 vLLM 端点上创建负载

向 vLLM 服务器创建负载，以测试 GKE 如何使用自定义 vLLM 指标进行自动扩缩。

  1. 运行 bash 脚本 (load.sh) 以向 vLLM 端点发送 N 个并发请求：

    #!/bin/bash
N=PARALLEL_PROCESSES
export vllm_service=$(kubectl get service vllm-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
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/Meta-Llama-3.1-70B", "prompt": "Write a story about san francisco", "max_tokens": 100, "temperature": 0}'
  done &  # Run in the background
done
wait

    将 PARALLEL_PROCESSES 替换为您要运行的并行进程数。

  2. 运行 bash 脚本：

    nohup ./load.sh &

验证 Google Cloud Managed Service for Prometheus 是否提取了指标

在 Google Cloud Managed Service for Prometheus 抓取指标并向 vLLM 端点添加负载后，您就可以在 Cloud Monitoring 上查看指标了。

  1. 在 Google Cloud 控制台中，前往 Metrics Explorer 页面。

    转到 Metrics Explorer

  2. 点击 < > PromQL。

  3. 输入以下查询，以观察流量指标：

    vllm:avg_generation_throughput_toks_per_s{cluster='CLUSTER_NAME_HERE'}

在折线图中，vLLM 指标从 0（加载前）扩大到某个值（加载后）。此图表可确认您的 vLLM 指标已被提取到 Google Cloud Managed Service for Prometheus。

下图是加载脚本执行后的图表示例。在本例中，模型服务器每秒提供约 2,000 个生成令牌。

测试

部署 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 部署 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

    GKE 会安排部署另一个 Pod，这会触发节点池自动扩缩器在部署第二个 vLLM 副本之前添加第二个节点。

  3. 观察 Pod 自动扩缩的进度：

    kubectl get hpa --watch

    输出类似于以下内容：

    NAME       REFERENCE             TARGETS       MINPODS   MAXPODS   REPLICAS   AGE
vllm-hpa   Deployment/vllm-tpu   <unknown>/5   1         2         0          6s
vllm-hpa   Deployment/vllm-tpu   34972m/5      1         2         1          16s
vllm-hpa   Deployment/vllm-tpu   25112m/5      1         2         2          31s
vllm-hpa   Deployment/vllm-tpu   35301m/5      1         2         2          46s
vllm-hpa   Deployment/vllm-tpu   25098m/5      1         2         2          62s
vllm-hpa   Deployment/vllm-tpu   35348m/5      1         2         2          77s

  4. 等待 10 分钟，然后重复验证 Google Cloud Managed Service for Prometheus 是否提取了指标部分中的步骤。Google Cloud Managed Service for Prometheus 会从两个 vLLM 端点提取指标。

test2

清理

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。

删除已部署的资源

为避免系统因您在本指南中创建的资源而向您的 Google Cloud 账号收取费用，请运行以下命令：

ps -ef | grep load.sh | awk '{print $2}' | xargs -n1 kill -9

gcloud container clusters delete CLUSTER_NAME \
  --location=ZONE

后续步骤