通过 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 可提高 TPU 上的服务吞吐量，具有以下功能：

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

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

支持的模型和功能

在 TPU 上使用 vLLM 部署 LLM 之前，我们建议您验证模型兼容性和支持的功能，以避免部署问题。The vllm-project/tpu-inference 库为 vLLM 在 TPU 上运行提供了必要的后端。

如需查看受支持模型和功能的完整列表，请参阅官方文档：

Cloud Storage FUSE

Cloud Storage FUSE 提供从您的 GKE 集群到 Cloud Storage 的访问权限，以获取驻留在对象存储桶中的模型权重。在本教程中，创建的 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. 部署模型

准备工作

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

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

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

  • Enable the required API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

  • 确保您对项目拥有以下一个或多个角色： roles/container.admin, roles/iam.serviceAccountAdmin, roles/iam.securityAdmin, roles/artifactregistry.writer, roles/container.clusterAdmin

    检查角色

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

      转到 IAM
    2. 选择项目。

    3. 在主账号 列中，找到标识您或您所属群组的所有行。如需了解您所属的群组，请与您的 管理员联系。

    4. 对于指定或包含您的所有行，请检查角色 列，查看 角色列表是否包含所需角色。

    授予角色

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

      转到 IAM
    2. 选择项目。
    3. 点击 授予访问权限。

    4. 在新的主账号 字段中，输入您的用户标识符。 这通常是员工身份池中的用户的标识符。如需了解详情， 请参阅 在 IAM 政策中表示员工池用户，或与您的管理员联系。

    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 && \
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：支持 TPU Trillium (v6e) 的 GKE 版本。如需了解详情，请参阅 验证 GKE 中的 TPU 可用性。
    • GSBUCKET：用于 Cloud Storage FUSE 的 Cloud Storage 存储桶的名称 。
    • KSA_NAME：用于访问 Cloud Storage 存储桶的 Kubernetes ServiceAccount 的名称。Cloud Storage FUSE 需要存储桶访问权限才能正常运行。
    • 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}

Standard

  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}. \
    --addons GcsFuseCsiDriver

此命令会创建一个 GKE Standard 集群， 并启用 工作负载身份联合 和 Cloud Storage FUSE CSI 驱动程序。如果您使用的是现有集群，请确保已启用 Cloud Storage FUSE CSI 驱动程序。

  1. 创建 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 Secret

  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 以访问存储桶

  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}./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}./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 文件缓存。

部署 vLLM 模型服务器

本教程将使用 Kubernetes Deployment 来部署 vLLM 模型服务器。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: vllm/vllm-tpu:latest
        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. 将 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}./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) 以向 vLLM 端点发送 N 个并发请求：

    #!/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 爬取指标并且您向 vLLM 端点添加负载后，您可以在 Cloud Monitoring 上查看指标。

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

    转到 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}

后续步骤