Esta página se ha traducido con Cloud Translation API.

Sirve un LLM mediante TPU Trillium en GKE con vLLM

En este tutorial se muestra cómo servir modelos de lenguaje extensos (LLMs) mediante unidades de procesamiento de tensor (TPUs) en Google Kubernetes Engine (GKE) con el framework de servicio vLLM. En este tutorial, se sirve Llama 3.1 70b, se usa TPU Trillium y se configura el autoescalado de pods horizontal mediante métricas del servidor vLLM.

Este documento es un buen punto de partida si necesitas el control granular, la escalabilidad, la resiliencia, la portabilidad y la rentabilidad de Kubernetes gestionado al desplegar y servir tus cargas de trabajo de IA o aprendizaje automático.

Fondo

Si usas TPU Trillium en GKE, puedes implementar una solución de servicio estable y lista para producción con todas las ventajas de Kubernetes gestionado, como una escalabilidad eficiente y una mayor disponibilidad. En esta sección se describen las tecnologías clave que se usan en esta guía.

TPU Trillium

Las TPUs son circuitos integrados para aplicaciones específicas (ASIC) desarrollados a medida por Google. Las TPUs se usan para acelerar los modelos de aprendizaje automático y de IA creados con frameworks como TensorFlow, PyTorch y JAX. En este tutorial se usa la TPU Trillium, que es la TPU de sexta generación de Google.

Antes de usar las TPUs en GKE, te recomendamos que completes el siguiente plan de formación:

Consulta información sobre la arquitectura del sistema de TPU Trillium.
Consulta información sobre las TPUs en GKE.

vLLM

vLLM es un framework de código abierto altamente optimizado para servir LLMs. vLLM puede aumentar el rendimiento del servicio en las TPUs con funciones como las siguientes:

Implementación optimizada de Transformer con PagedAttention.
Agrupación continua para mejorar el rendimiento general del servicio.
Paralelismo de tensores y servicio distribuido en varias TPUs.

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

Nota: Este tutorial se centra en el despliegue de vLLM en una configuración de un solo host, que es ideal para modelos que se pueden servir desde un solo slice de TPU, como Llama 3.1 70b en un tipo de máquina ct6e-standard-8t. Es importante tener en cuenta que las configuraciones de varios hosts no se admiten cuando se usa vLLM con TPUs en GKE. La falta de compatibilidad con varios hosts limita el uso de vLLM para proporcionar modelos extremadamente grandes (por ejemplo, con más de 400.000 millones de parámetros) que requieren la memoria y la computación agregadas de varios hosts. En el caso de los sistemas o modelos de producción que requieren una configuración multihost, la solución recomendada y optimizada para el rendimiento es usar JetStream, el motor de Google para la inferencia de TPU. Para empezar a usar una implementación de varios hosts, consulta Inferencia de MaxText de JetStream en TPU v6e.

FUSE de Cloud Storage

Cloud Storage FUSE proporciona acceso desde tu clúster de GKE a Cloud Storage para los pesos del modelo que residen en segmentos de almacenamiento de objetos. En este tutorial, el segmento 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 segmento de Cloud Storage. Cuando se reinicie un pod o se aumente la escala de una implementación, las cargas de modelos posteriores descargarán los datos almacenados en caché del segmento de Cloud Storage, lo que aprovechará las descargas paralelas para optimizar el rendimiento.

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

Crear un clúster de GKE

Puedes servir LLMs en TPUs en un clúster Autopilot o Estándar de GKE. Te recomendamos que uses un clúster de Autopilot para disfrutar de una experiencia de Kubernetes totalmente gestionada. Para elegir el modo de funcionamiento de GKE que mejor se adapte a tus cargas de trabajo, consulta Elegir un modo de funcionamiento de GKE.

Autopilot

Crea un clúster de Autopilot de GKE:

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

Estándar

Crea un clúster estándar de GKE:

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 segmento 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 Estándar de GKE que use Workload Identity Federation for GKE y tenga habilitado el controlador CSI de Cloud Storage FUSE.
- Un grupo de nodos de TPU Trillium con un tipo de máquina ct6e-standard-8t. Este grupo de nodos tiene un nodo, ocho chips de TPU y el autoescalado habilitado.

Configurar kubectl para que se comunique con el clúster

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

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

Crear un secreto de Kubernetes para las credenciales de Hugging Face

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

Crea un secreto de Kubernetes que contenga el token de Hugging Face. Para ello, ejecuta el siguiente comando:

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

Crea un segmento de Cloud Storage

En Cloud Shell, ejecuta el siguiente comando:

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

De esta forma, se crea un segmento de Cloud Storage para almacenar los archivos del modelo que descargues de Hugging Face.

Configurar una cuenta de servicio de Kubernetes para acceder al segmento

Crea la cuenta de servicio de Kubernetes:

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

Concede acceso de lectura y escritura a la cuenta de servicio de Kubernetes para acceder al segmento 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"

También puedes conceder acceso de lectura y escritura a todos los segmentos de Cloud Storage del 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 segmento de Cloud Storage para almacenar el modelo descargado y la caché de compilación. Un controlador CSI de FUSE de Cloud Storage lee el contenido del segmento.
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 en función del tamaño esperado del contenido del modelo (por ejemplo, los archivos de pesos). En este tutorial, se usa la caché de archivos de Cloud Storage FUSE respaldada por RAM.

Desplegar el servidor de modelos vLLM

Para desplegar el servidor de modelos vLLM, en este tutorial se usa un despliegue de Kubernetes. Un Deployment es un objeto de la API de Kubernetes que te permite ejecutar varias réplicas de pods distribuidas entre los nodos de un clúster.

Inspecciona el siguiente archivo de 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 el número de réplicas de la implementación, 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:

Elimina la ubicación de la caché de XLA quitando 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"
```
Escala la implementación a 1 y espera a que la primera réplica esté lista y escriba en la caché de XLA. A continuación, escala a réplicas adicionales. De esta forma, el resto de las réplicas pueden leer la caché sin intentar escribir en ella.

Aplica el manifiesto ejecutando el siguiente comando:

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

Consulta los registros del servidor del modelo en ejecución:

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

La salida debería ser similar a la 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)

Aplicar el modelo

Para obtener la dirección IP externa del servicio 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 mediante 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
}'

La salida debería ser similar a la siguiente:

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

Configurar el escalador automático personalizado

En esta sección, configurará el autoescalado horizontal de pods mediante métricas de Prometheus personalizadas. Usas las métricas de Google Cloud Managed Service para Prometheus del servidor 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 Stackdriver de métricas personalizadas 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

Añade el rol Lector 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 Escritor de métricas de monitorización. En este tutorial se usa la cuenta de servicio predeterminada de Compute Engine.

Guarda el siguiente archivo de 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}

Crear carga en el endpoint de vLLM

Crea carga en el servidor vLLM para probar cómo se adapta dinámicamente GKE con una métrica vLLM personalizada.

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

Sustituye PARALLEL_PROCESSES por el número de procesos paralelos que quieras ejecutar.

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

Verificar que Google Cloud Managed Service para Prometheus ingiere las métricas

Una vez que Managed Service para Prometheus de Google Cloud haya recogido las métricas y añadas carga al endpoint de vLLM, podrás ver las métricas en Cloud Monitoring.

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

Ir a Explorador de métricas
Haz clic en < > PromQL.
Introduce 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 vLLM aumenta de 0 (precarga) a un valor (después de la carga). Este gráfico confirma que las métricas de tu vLLM se están ingiriendo en 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 máximo posterior a la carga de casi 400 en un minuto.

vllm:num_requests_waiting

Desplegar la configuración de la herramienta de adaptación dinámica horizontal de pods

A la hora de decidir en qué métrica basar el autoescalado, te recomendamos las siguientes métricas para la TPU de vLLM:

num_requests_waiting: esta métrica se refiere al número de solicitudes que esperan en la cola del servidor del modelo. Este número empieza a crecer de forma notable cuando la caché de pares clave-valor está llena.
gpu_cache_usage_perc: esta métrica está relacionada con la utilización de la caché de clave-valor, que se corresponde directamente con el número de solicitudes que se procesan en un ciclo de inferencia determinado en el servidor del modelo. Ten en cuenta que esta métrica funciona igual en las GPUs y las TPUs, aunque está vinculada al esquema de nomenclatura de las GPUs.

Te recomendamos que uses num_requests_waiting cuando optimices el rendimiento y el coste, y cuando puedas alcanzar tus objetivos de latencia con el rendimiento máximo del servidor de tu modelo.

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

Para obtener más información, consulta el artículo Prácticas recomendadas para escalar automáticamente cargas de trabajo de inferencia de modelos de lenguaje extensos (LLMs) con TPUs.

Al seleccionar un averageValueobjetivo para tu configuración de HPA, tendrás que determinarlo de forma experimental. Consulta la entrada de blog Ahorra en GPUs: autoescalado 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 con vLLM TPU.

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

Guarda el siguiente archivo de 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 de Google Cloud Managed Service para Prometheus siguen el formato vllm:metric_name.

Práctica recomendada:

Usa num_requests_waiting para escalar el rendimiento. Usa gpu_cache_usage_perc para los casos prácticos de TPU sensibles a la latencia.

Despliega la configuración de la herramienta de adaptación dinámica horizontal de pods:
```
kubectl apply -f vllm-hpa.yaml -n ${NAMESPACE}
```
GKE programa otro pod para desplegarlo, lo que activa el escalador automático del grupo de nodos para añadir un segundo nodo antes de desplegar la segunda réplica de vLLM.

Supervisa el progreso del autoescalado de pods:

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

El resultado debería ser similar al siguiente:

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 Verificar que Google Cloud Managed Service para Prometheus ingiere las métricas. Google Cloud Managed Service para Prometheus ahora ingiere las métricas de ambos endpoints de vLLM.