通过 vLLM，使用 GKE 中的 TPU Trillium 提供 LLM

Standard Autopilot

本教程介绍如何通过 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 之前，我们建议您完成以下学习路线：

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

vLLM

vLLM 是一个经过高度优化的开源 LLM 部署框架。vLLM 可提高 TPU 上的服务吞吐量，具有以下功能：

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

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

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 工程师或平台管理员。

本教程介绍以下步骤：

根据模型特征创建一个具有推荐 TPU Trillium 拓扑的 GKE 集群。
在集群中的节点池上部署 vLLM 框架。
使用 vLLM 框架通过负载均衡器部署 Llama 3.1 70b。
使用 vLLM 服务器指标设置 Pod 横向自动扩缩。
部署模型

准备工作

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 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.
  Go to IAM
2. 选择项目。
3. 点击 授予访问权限。
4. 在新的主账号字段中，输入您的用户标识符。这通常是 Google 账号的电子邮件地址。
5. 在选择角色列表中，选择一个角色。
6. 如需授予其他角色，请点击 添加其他角色，然后添加其他各个角色。
7. 点击 Save（保存）。

如果您还没有 Hugging Face 账号，请创建一个。
确保您的项目具有足够的配额，以便用于 GKE 中的 Cloud TPU。

准备环境

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

获取对模型的访问权限

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

生成一个访问令牌

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

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

启动 Cloud Shell

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

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

在 Google Cloud 控制台中，点击 Google Cloud 控制台中的 激活 Cloud Shell 以启动 Cloud Shell 会话。此操作会在 Google Cloud 控制台的底部窗格中启动会话。
设置默认环境变量：
```
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 ZONE=ZONE && \
export REGION=REGION && \
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 集群的名称。
- ZONE：支持 TPU Trillium (v6e) 的可用区。
- REGION：支持 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 存储桶的 Kubernetes ServiceAccount 的名称。Cloud Storage FUSE 需要存储桶访问权限才能正常运行。
- NAMESPACE：您要部署 vLLM 资产的 Kubernetes 命名空间。

创建 GKE 集群

您可以在 GKE Autopilot 或 Standard 集群中的 TPU 上部署 LLM。我们建议您使用 Autopilot 集群获得全托管式 Kubernetes 体验。如需选择最适合您的工作负载的 GKE 操作模式，请参阅选择 GKE 操作模式。

Autopilot

创建 GKE Autopilot 集群：

gcloud container clusters create-auto ${CLUSTER_NAME} \
    --cluster-version=${CLUSTER_VERSION} \
    --region=${REGION}

Standard

创建 GKE Standard 集群：

gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --region=${REGION} \
    --node-locations=${ZONE} \
    --cluster-version=${CLUSTER_VERSION} \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --addons GcsFuseCsiDriver

创建 TPU 切片节点池：
```
gcloud container node-pools create tpunodepool \
    --region=${REGION} \
    --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 创建以下资源：
- 使用 Workload Identity Federation for GKE 并已启用 Cloud Storage FUSE CSI 驱动程序的 GKE Standard 集群。
- 具有 ct6e-standard-8t 机器类型的 TPU Trillium 节点池。此节点池有一个节点、八个 TPU 芯片，并启用了自动扩缩。

配置 kubectl 以与您的集群通信

如需配置 kubectl 以与您的集群通信，请运行以下命令：

  gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}

为 Hugging Face 凭据创建 Kubernetes Secret

创建命名空间。如果您使用的是 default 命名空间，则可以跳过此步骤：
```
kubectl create namespace ${NAMESPACE}
```

创建包含 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 以访问存储桶

创建 Kubernetes ServiceAccount：

kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}

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

或者，您可以授予对项目中所有 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 文件缓存。

部署 vLLM 模型服务器

本教程将使用 Kubernetes Deployment 来部署 vLLM 模型服务器。Deployment 是一个 Kubernetes API 对象，可让您运行在集群节点中分布的多个 Pod 副本。

检查保存为 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

通过运行以下命令来应用清单：

kubectl apply -f vllm-llama3-70b.yaml -n ${NAMESPACE}

查看正在运行的模型服务器的日志：

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)

应用模型

如需获取 VLLM 服务的外部 IP 地址，请运行以下命令：

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

使用 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 集群上默认启用。

对集群设置自定义指标 Stackdriver 适配器：

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml

将 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

将以下清单保存为 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

将其应用于集群：

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

在 vLLM 端点上创建负载

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

运行 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 替换为您要运行的并行进程数。

运行 bash 脚本：
```
chmod +x load.sh
nohup ./load.sh &
```

验证 Google Cloud Managed Service for Prometheus 是否注入指标

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

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

转到 Metrics Explorer
点击 < > PromQL。

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

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 横向自动扩缩器配置，请按照以下步骤操作：

将以下清单保存为 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。

部署 Pod 横向自动扩缩器配置：
```
kubectl apply -f vllm-hpa.yaml -n ${NAMESPACE}
```
GKE 会调度另一个 Pod 进行部署，这会触发节点池自动扩缩器在部署第二个 vLLM 副本之前添加第二个节点。

观察 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

等待 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} \
  --region=${REGION}

通过 vLLM，使用 GKE 中的 TPU Trillium 提供 LLM 使用集合让一切井井有条 根据您的偏好保存内容并对其进行分类。

背景

TPU Trillium

vLLM

Cloud Storage FUSE

目标

准备工作

Check for the roles

Grant the roles

准备环境

获取对模型的访问权限

生成一个访问令牌

启动 Cloud Shell

创建 GKE 集群

Autopilot

Standard

配置 kubectl 以与您的集群通信

为 Hugging Face 凭据创建 Kubernetes Secret

创建 Cloud Storage 存储桶

设置 Kubernetes ServiceAccount 以访问存储桶

部署 vLLM 模型服务器

应用模型

设置自定义自动扩缩器

在 vLLM 端点上创建负载

验证 Google Cloud Managed Service for Prometheus 是否注入指标

部署 Pod 横向自动扩缩器配置

清理

删除已部署的资源

后续步骤

通过 vLLM，使用 GKE 中的 TPU Trillium 提供 LLM