Se usó la API de Cloud Translation para traducir esta página.

Entrega un LLM con TPU Trillium en GKE con vLLM

En este instructivo, se muestra cómo entregar modelos de lenguaje grandes (LLM) con unidades de procesamiento tensorial (TPU) en Google Kubernetes Engine (GKE) con el framework de entrega vLLM. En este instructivo, entregarás Llama 3.1 70b, usarás TPU Trillium y configurarás el escalamiento automático horizontal de Pods con las métricas del servidor de vLLM.

Este documento es un buen punto de partida si necesitas el control detallado, la escalabilidad, la resiliencia, la portabilidad y la rentabilidad de Kubernetes administrado cuando implementas y entregas tus cargas de trabajo de IA/AA.

Fondo

Con la TPU Trillium en GKE, puedes implementar una solución de entrega sólida y lista para la producción con todos los beneficios de administrar Kubernetes, incluida la escalabilidad eficiente y la mayor disponibilidad. En esta sección, se describen las tecnologías clave que se usan en este instructivo.

TPU Trillium

Las TPU son circuitos integrados específicos de aplicaciones (ASIC) desarrollados de forma personalizada por Google. Las TPU se usan para acelerar el aprendizaje automático y los modelos de IA compilados con frameworks como TensorFlow, PyTorch y JAX. En este instructivo, se usa la TPU Trillium, que es la sexta generación de TPU de Google.

Antes de usar las TPU en GKE, te recomendamos que completes la siguiente ruta de aprendizaje:

Obtén información sobre la arquitectura del sistema de la TPU Trillium.
Obtén información sobre las TPU en GKE.

vLLM

vLLM es un framework de código abierto altamente optimizado para la entrega de LLMs. vLLM puede aumentar la capacidad de procesamiento de entrega en TPUs, con funciones como las siguientes:

Implementación optimizada de transformadores con PagedAttention.
Agrupación en lotes continua para mejorar la capacidad de procesamiento general de la entrega
Paralelismo de tensor y entrega distribuida en varias TPU

Para obtener más información, consulta la documentación de vLLM.

Nota: Este instructivo se enfoca en la implementación de vLLM en una configuración de un solo host, que es ideal para modelos que se pueden entregar desde un solo segmento de TPU, como Llama 3.1 70b en un tipo de máquina ct6e-standard-8t. Es importante tener en cuenta que no se admiten las configuraciones de varios hosts cuando se usa vLLM con TPU en GKE. La falta de compatibilidad con varios hosts limita el uso de vLLM para entregar modelos extremadamente grandes (por ejemplo, más de 400,000 millones de parámetros) que requieren la memoria y el procesamiento agregados de varios hosts. Para los sistemas o modelos de producción que requieren una configuración de varios hosts, la solución recomendada y optimizada para el rendimiento es usar JetStream, el motor de Google para la inferencia en TPU. Para comenzar con una implementación de varios hosts, consulta Inferencia de JetStream MaxText en la TPU v6e.

Cloud Storage FUSE

Cloud Storage FUSE proporciona acceso desde tu clúster de GKE a Cloud Storage para los pesos del modelo que residen en buckets de almacenamiento de objetos. En este instructivo, el bucket de Cloud Storage creado estará vacío al principio. Cuando se inicia vLLM, GKE descarga el modelo de Hugging Face y almacena en caché los pesos en el bucket de Cloud Storage. Cuando se reinicia el Pod o se aumenta la escala de la implementación, las cargas posteriores del modelo descargarán los datos almacenados en caché del bucket de Cloud Storage, lo que aprovechará las descargas paralelas para obtener un rendimiento óptimo.

Para obtener más información, consulta la documentación del controlador de CSI de Cloud Storage FUSE.

Crea un clúster de GKE

Puedes entregar LLM en TPU en un clúster de GKE Autopilot o Standard. Te recomendamos que uses un clúster de Autopilot para una experiencia de Kubernetes completamente administrada. Para elegir el modo de operación de GKE que se adapte mejor a tus cargas de trabajo, consulta Elige un modo de operación de GKE.

Autopilot

Crea un clúster de GKE Autopilot:

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

Estándar

Crea un clúster de 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

Crea un grupo de nodos de porción de 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 crea los siguientes recursos para el LLM:
- Un clúster de GKE Standard que usa la federación de identidades para cargas de trabajo para GKE y tiene habilitado el controlador de CSI de Cloud Storage FUSE.
- Un grupo de nodos TPU Trillium con un tipo de máquina ct6e-standard-8t. Este grupo de nodos tiene un nodo, ocho chips TPU y el ajuste de escala automático habilitado.

Configura kubectl para comunicarse con tu clúster

Para configurar kubectl para que se comunique con tu clúster, ejecuta el siguiente comando:

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

Crea un secreto de Kubernetes para las credenciales de Hugging Face

Crea un espacio de nombres. Puedes omitir este paso si usas el espacio de nombres default:
```
kubectl create namespace ${NAMESPACE}
```

Ejecuta el siguiente comando para crear un Secret de Kubernetes que contenga el token de Hugging Face:

kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --namespace ${NAMESPACE}

Cree un bucket de Cloud Storage

En Cloud Shell, ejecute el siguiente comando:

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

Esto crea un bucket de Cloud Storage para almacenar los archivos del modelo que descargas de Hugging Face.

Configura una ServiceAccount de Kubernetes para acceder al bucket

Crea la ServiceAccount de Kubernetes:

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

Otorga acceso de lectura y escritura a la ServiceAccount de Kubernetes para acceder al bucket de 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"

Como alternativa, puedes otorgar acceso de lectura y escritura a todos los buckets de Cloud Storage en el proyecto:
```
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 crea los siguientes recursos para el LLM:
1. Un bucket de Cloud Storage para almacenar el modelo descargado y la caché de compilación Un controlador de CSI de Cloud Storage FUSE lee el contenido del bucket.
2. Volúmenes con el almacenamiento en caché de archivos habilitado y la función de descarga paralela de Cloud Storage FUSE
Práctica recomendada:
Usa una caché de archivos respaldada por tmpfs o Hyperdisk / Persistent Disk, según el tamaño esperado del contenido del modelo, por ejemplo, los archivos de pesos. En este instructivo, usarás la caché de archivos de Cloud Storage FUSE respaldada por RAM.

Implementa el servidor de modelos de vLLM

Para implementar el servidor del modelo de vLLM, en este instructivo, se usa una implementación de Kubernetes. Un Deployment es un objeto de la API de Kubernetes que te permite ejecutar varias réplicas de Pods que se distribuyen entre los nodos de un clúster.

Inspecciona el siguiente manifiesto de Deployment guardado como vllm-llama3-70b.yaml, que usa una sola réplica:


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

Si aumentas la cantidad de réplicas de la Deployment, las escrituras simultáneas en VLLM_XLA_CACHE_PATH provocarán el error RuntimeError: filesystem error: cannot create directories. Para evitar este error, tienes dos opciones:

Para quitar la ubicación de la caché de XLA, quita el siguiente bloque del archivo YAML de implementación. Esto significa que todas las réplicas volverán a compilar la caché.
```
- name: VLLM_XLA_CACHE_PATH
  value: "/data"
```
Ajusta la escala de la Deployment a 1 y espera a que la primera réplica esté lista y escriba en la caché de XLA. Luego, escala a réplicas adicionales. Esto permite que el resto de las réplicas lean la caché sin intentar escribirla.

Aplica el manifiesto ejecutando el siguiente comando:

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

Observa los registros del servidor de modelos en ejecución:

kubectl logs -f -l app=vllm-tpu -n ${NAMESPACE}

El resultado debería ser similar al siguiente:

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)

Entrega el modelo

Para obtener la dirección IP externa del servicio de VLLM, ejecuta el siguiente comando:

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

Interactúa con el modelo usando 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
}'

El resultado debería ser similar al siguiente ejemplo:

{"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}}

Configura el escalador automático personalizado

En esta sección, configurarás el ajuste de escala automático horizontal de Pods con métricas personalizadas de Prometheus. Usas las métricas de Google Cloud Managed Service para Prometheus del servidor de vLLM.

Para obtener más información, consulta Google Cloud Managed Service para Prometheus. Esta opción debería estar habilitada de forma predeterminada en el clúster de GKE.

Configura el adaptador de métricas personalizadas de Stackdriver en tu clúster:

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

Agrega el rol de visualizador de Monitoring a la cuenta de servicio que usa el adaptador de Stackdriver de métricas personalizadas:
```
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
```
Nota: Asegúrate de que la cuenta de servicio que usa tu clúster de GKE tenga el rol de Escritor de métricas de Monitoring. En este instructivo, se usa la cuenta de servicio predeterminada de Compute Engine.

Guarda el siguiente manifiesto como 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

Aplícalo al clúster:

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

Genera carga en el extremo de vLLM

Crea carga en el servidor de vLLM para probar cómo GKE ajusta la escala automáticamente con una métrica personalizada de vLLM.

Ejecuta una secuencia de comandos de Bash (load.sh) para enviar N cantidad de solicitudes paralelas al extremo de 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

Reemplaza PARALLEL_PROCESSES por la cantidad de procesos paralelos que deseas ejecutar.

Ejecuta la secuencia de comandos de Bash:
```
chmod +x load.sh
nohup ./load.sh &
```

Verifica que Google Cloud Managed Service para Prometheus transfiera las métricas

Después de que Google Cloud Managed Service para Prometheus extraiga las métricas y agregues carga al extremo de vLLM, podrás ver las métricas en Cloud Monitoring.

En la consola de Google Cloud , ve a la página Explorador de métricas.

Ir al Explorador de métricas
Haz clic en < > PromQL.
Ingresa la siguiente consulta para observar las métricas de tráfico:
```
vllm:num_requests_waiting{cluster='CLUSTER_NAME'}
```

Un gráfico de líneas muestra tu métrica de vLLM (num_requests_waiting) medida a lo largo del tiempo. La métrica de vLLM se incrementa de 0 (antes de la carga) a un valor (después de la carga). Este gráfico confirma que tus métricas de vLLM se transfieren a Google Cloud Managed Service para Prometheus. En el siguiente gráfico de ejemplo, se muestra un valor de precarga inicial de 0, que alcanza un valor de poscarga máximo de casi 400 en un minuto.

vllm:num_requests_waiting

Implementa la configuración del Horizontal Pod Autoscaler

Cuando decidas en qué métrica basar el ajuste de escala automático, te recomendamos las siguientes métricas para la TPU de vLLM:

num_requests_waiting: Esta métrica se relaciona con la cantidad de solicitudes en espera en la cola del servidor del modelo. Este número comienza a aumentar de forma notable cuando la caché de KV está llena.
gpu_cache_usage_perc: Esta métrica se relaciona con la utilización de la caché de KV, que se correlaciona directamente con la cantidad de solicitudes que se procesan para un ciclo de inferencia determinado en el servidor del modelo. Ten en cuenta que esta métrica funciona de la misma manera en las GPU y las TPU, aunque está vinculada al esquema de nomenclatura de las GPU.

Te recomendamos que uses num_requests_waiting cuando optimices la capacidad de procesamiento y el costo, y cuando tus objetivos de latencia se puedan alcanzar con la capacidad de procesamiento máxima del servidor de modelos.

Te recomendamos que uses gpu_cache_usage_perc cuando tengas cargas de trabajo sensibles a la latencia en las que el escalamiento basado en colas no sea lo suficientemente rápido para cumplir con tus requisitos.

Para obtener más información, consulta Prácticas recomendadas para el ajuste de escala automático de cargas de trabajo de inferencia de modelos de lenguaje grandes (LLM) con TPU.

Cuando selecciones un averageValue objetivo para tu configuración de HPA, deberás determinarlo de forma experimental. Consulta la entrada de blog Ahorra en GPUs: Ajuste de escala más inteligente para tus cargas de trabajo de inferencia de GKE para obtener más ideas sobre cómo optimizar esta parte. El profile-generator que se usa en esta entrada de blog también funciona para la TPU de vLLM.

En las siguientes instrucciones, implementarás tu configuración de HPA con la métrica num_requests_waiting. Para fines de demostración, establece la métrica en un valor bajo para que la configuración del HPA escale tus réplicas de vLLM a dos. Para implementar la configuración del Horizontal Pod Autoscaler con num_requests_waiting, sigue estos pasos:

Guarda el siguiente manifiesto como 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

Las métricas de vLLM en Google Cloud Managed Service para Prometheus siguen el formato vllm:metric_name.

Práctica recomendada:

Usa num_requests_waiting para ajustar la capacidad de procesamiento. Usa gpu_cache_usage_perc para los casos de uso de TPU sensibles a la latencia.

Implementa la configuración del Horizontal Pod Autoscaler:
```
kubectl apply -f vllm-hpa.yaml -n ${NAMESPACE}
```
GKE programa otro Pod para implementar, lo que activa el escalador automático del grupo de nodos para agregar un segundo nodo antes de implementar la segunda réplica de vLLM.

Observa el progreso del ajuste de escala automático de Pods:

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

El resultado es similar a este:

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

Espera 10 minutos y repite los pasos de la sección Verifica que Google Cloud Managed Service para Prometheus ingiere las métricas. Google Cloud Managed Service para Prometheus ahora admite las métricas de ambos extremos de vLLM.