Entrega un LLM (Llama3.1 405B) con varios nodos de GPU


Descripción general

En este instructivo, se muestra cómo entregar Llama 3.1 405b con unidades de procesamiento gráfico (GPU) en varios nodos en Google Kubernetes Engine (GKE) con el framework de entrega vLLM y la API de LeaderWorkerSet (LWS).

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.

LeaderWorkerSet (LWS)

LWS es una API de implementación de Kubernetes que aborda patrones de implementación comunes de cargas de trabajo de inferencia de varios nodos de IA/AA. LWS permite tratar varios pods como un grupo.

Publicación de varios hosts con vLLM

Cuando implementes modelos de lenguaje excepcionalmente grandes que no se ajusten a un solo nodo de GPU, usa varios nodos de GPU para entregar el modelo. vLLM admite el paralelismo de tensores y el paralelismo de canalizaciones para ejecutar cargas de trabajo en varias GPUs.

El paralelismo de tensor divide las multiplicaciones de matrices en la capa de transformador en varias GPUs. Sin embargo, esta estrategia requiere una red rápida debido a la comunicación necesaria entre las GPUs, lo que la hace menos adecuada para ejecutar cargas de trabajo en varios nodos.

El paralelismo de canalización divide el modelo por capa o verticalmente. Esta estrategia no requiere una comunicación constante entre las GPUs, lo que la convierte en una mejor opción cuando se ejecutan modelos en varios nodos.

Puedes usar ambas estrategias en la publicación de varios nodos. Por ejemplo, cuando usas dos nodos con 8 GPUs H100 cada uno, puedes usar el paralelismo de canalización de dos vías para dividir el modelo entre los dos nodos y el paralelismo de tensor de ocho vías para dividir el modelo entre las ocho GPUs de cada nodo.

Objetivos

  1. Prepara un clúster de GKE Standard.
  2. Implementa vLLM en varios nodos de tu clúster.
  3. Usa vLLM para entregar el modelo Llama3 405b a través de curl.

Antes de comenzar

  • 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

    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.

      Ir a IAM
    2. Selecciona el proyecto.
    3. Haz clic en Grant access.
    4. En el campo Principales nuevas, ingresa tu identificador de usuario. Esta suele ser la dirección de correo electrónico de una Cuenta de Google.

    5. En la lista Seleccionar un rol, elige un rol.
    6. Para otorgar funciones adicionales, haz clic en Agregar otro rol y agrega cada rol adicional.
    7. Haz clic en Guardar.

Obtén acceso al modelo

Genera un token de acceso

Si aún no tienes uno, genera un nuevo token de Hugging Face:

  1. Haz clic en Tu perfil > Configuración > Tokens de acceso.
  2. Selecciona Token nuevo.
  3. Especifica el nombre que desees y un rol de al menos Read.
  4. Selecciona Generate un token.

Prepare el entorno

En este instructivo, usarás Cloud Shell para administrar recursos alojados en Google Cloud. Cloud Shell tiene preinstalado el software que necesitarás para este instructivo, incluidos kubectl y la CLI de gcloud.

Para configurar tu entorno con Cloud Shell, sigue estos pasos:

  1. En la consola de Google Cloud , haz clic en Ícono de activación de Cloud Shell Activar Cloud Shell en la consola de Google Cloud para iniciar una sesión de Cloud Shell. Esto inicia una sesión en el panel inferior de la consola de Google Cloud .

  2. Configura las variables de entorno predeterminadas:

    gcloud config set project PROJECT_ID
    export PROJECT_ID=$(gcloud config get project)
    export CLUSTER_NAME=CLUSTER_NAME
    export ZONE=ZONE
    export HF_TOKEN=HUGGING_FACE_TOKEN
    export IMAGE_NAME=IMAGE_NAME
    

    Reemplaza los siguientes valores:

    • PROJECT_ID: El ID del proyecto de Google Cloud .
    • CLUSTER_NAME: Es el nombre del clúster de GKE.
    • ZONE: Es una zona que admite H100.
    • IMAGE_NAME: La imagen de vLLM que incluye la secuencia de comandos de Ray. Te proporcionamos us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240821_1034_RC00, o bien puedes crear el tuyo.

Cree un clúster de GKE

Puedes entregar modelos con vLLM en varios nodos de GPU 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

En Cloud Shell, ejecuta el siguiente comando:

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

Estándar

  1. Crea un clúster de GKE Standard con dos nodos de CPU:

    gcloud container clusters create CLUSTER_NAME \
        --project=PROJECT_ID \
        --num-nodes=2 \
        --location=ZONE \
        --machine-type=e2-standard-16
    
  2. Crea un grupo de nodos A3 con dos nodos, cada uno con ocho H100:

    gcloud container node-pools create gpu-nodepool \
        --location=ZONE \
        --num-nodes=2 \
        --machine-type=a3-highgpu-8g \
        --accelerator=type=nvidia-h100-80gb,count=8,gpu-driver-version=LATEST \
        --placement-type=COMPACT \
        --cluster=CLUSTER_NAME
    

Configura kubectl para comunicarse con tu clúster:

gcloud container clusters get-credentials CLUSTER_NAME --location=ZONE

Crea un secreto de Kubernetes para las credenciales de Hugging Face

Crea un Secret de Kubernetes que contenga el token de Hugging Face:

kubectl create secret generic hf-secret \
  --from-literal=hf_api_token=${HF_TOKEN} \
  --dry-run=client -o yaml | kubectl apply -f -

(Opcional) Crea tu propia imagen de vLLM de varios nodos

Elige esta opción para tener más control sobre el contenido de tu imagen de Docker y para incluir dependencias específicas junto con tu secuencia de comandos. Para ejecutar vLLM en varios nodos, puedes usar Ray para la comunicación entre nodos. Puedes ver el Dockerfile, que contiene una secuencia de comandos de Bash para configurar Ray con vLLM en el repositorio de LeaderWorkerSet.

Compile el contenedor

  1. Clona el repositorio LeaderWorkerSet:

    git clone https://github.com/kubernetes-sigs/lws.git
    
  2. Compila la imagen.

    cd lws/docs/examples/vllm/build/ && docker build -f Dockerfile . -t vllm-multihost
    

Envía la imagen al registro de Artifact Registry.

Para asegurarte de que tu implementación de Kubernetes pueda acceder a la imagen, almacénala en Artifact Registry dentro de tu proyecto de Google Cloud .

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

Instala LeaderWorkerSet

Para instalar LWS, ejecuta el siguiente comando:

VERSION=v0.4.2
kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/$VERSION/manifests.yaml

Valida que el controlador LeaderWorkerSet se ejecute en el espacio de nombres lws-system:

kubectl get pod -n lws-system

El resultado es similar a este:

NAME                                      READY   STATUS    RESTARTS   AGE
lws-controller-manager-5c4ff67cbd-9jsfc   2/2     Running   0          6d23h

Implementa el servidor de modelos de vLLM

Para implementar el servidor de modelos de vLLM, sigue estos pasos:

  1. Inspecciona el manifiesto vllm-llama3-405b-A3.yaml.

    
    apiVersion: leaderworkerset.x-k8s.io/v1
    kind: LeaderWorkerSet
    metadata:
      name: vllm
    spec:
      replicas: 1
      leaderWorkerTemplate:
        size: 2
        restartPolicy: RecreateGroupOnPodRestart
        leaderTemplate:
          metadata:
            labels:
              role: leader
          spec:
            nodeSelector:
              cloud.google.com/gke-accelerator: nvidia-h100-80gb
            containers:
              - name: vllm-leader
                image: IMAGE_NAME
                env:
                  - name: HUGGING_FACE_HUB_TOKEN
                    valueFrom:
                      secretKeyRef:
                        name: hf-secret
                        key: hf_api_token
                command:
                  - sh
                  - -c
                  - "/workspace/vllm/examples/ray_init.sh leader --ray_cluster_size=$(LWS_GROUP_SIZE); 
                    python3 -m vllm.entrypoints.openai.api_server --port 8080 --model meta-llama/Meta-Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size 2"
                resources:
                  limits:
                    nvidia.com/gpu: "8"
                ports:
                  - containerPort: 8080
                readinessProbe:
                  tcpSocket:
                    port: 8080
                  initialDelaySeconds: 15
                  periodSeconds: 10
                volumeMounts:
                  - mountPath: /dev/shm
                    name: dshm
            volumes:
            - name: dshm
              emptyDir:
                medium: Memory
                sizeLimit: 15Gi
        workerTemplate:
          spec:
            nodeSelector:
              cloud.google.com/gke-accelerator: nvidia-h100-80gb
            containers:
              - name: vllm-worker
                image: IMAGE_NAME
                command:
                  - sh
                  - -c
                  - "/workspace/vllm/examples/ray_init.sh worker --ray_address=$(LWS_LEADER_ADDRESS)"
                resources:
                  limits:
                    nvidia.com/gpu: "8"
                env:
                  - name: HUGGING_FACE_HUB_TOKEN
                    valueFrom:
                      secretKeyRef:
                        name: hf-secret
                        key: hf_api_token
                volumeMounts:
                  - mountPath: /dev/shm
                    name: dshm   
            volumes:
            - name: dshm
              emptyDir:
                medium: Memory
                sizeLimit: 15Gi
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: vllm-leader
    spec:
      ports:
        - name: http
          port: 8080
          protocol: TCP
          targetPort: 8080
      selector:
        leaderworkerset.sigs.k8s.io/name: vllm
        role: leader
      type: ClusterIP
    
  2. Ejecuta el siguiente comando para aplicar el manifiesto:

    kubectl apply -f vllm-llama3-405b-A3.yaml
    
  3. Consulta los registros del servidor de modelos en ejecución

    kubectl logs vllm-0 -c vllm-leader
    

    El resultado debería ser similar al siguiente:

    INFO 08-09 21:01:34 api_server.py:297] Route: /detokenize, Methods: POST
    INFO 08-09 21:01:34 api_server.py:297] Route: /v1/models, Methods: GET
    INFO 08-09 21:01:34 api_server.py:297] Route: /version, Methods: GET
    INFO 08-09 21:01:34 api_server.py:297] Route: /v1/chat/completions, Methods: POST
    INFO 08-09 21:01:34 api_server.py:297] Route: /v1/completions, Methods: POST
    INFO 08-09 21:01:34 api_server.py:297] Route: /v1/embeddings, Methods: POST
    INFO:     Started server process [7428]
    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

Ejecuta el siguiente comando para configurar la redirección de puertos al modelo:

kubectl port-forward svc/vllm-leader 8080:8080

Interactúa con el modelo con curl

En una terminal nueva, envía una solicitud al servidor:

curl http://localhost:8080/v1/completions \
-H "Content-Type: application/json" \
-d '{
    "model": "meta-llama/Meta-Llama-3.1-405B-Instruct",
    "prompt": "San Francisco is a",
    "max_tokens": 7,
    "temperature": 0
}'

El resultado debería ser similar al siguiente ejemplo:

{"id":"cmpl-0a2310f30ac3454aa7f2c5bb6a292e6c",
"object":"text_completion","created":1723238375,"model":"meta-llama/Meta-Llama-3.1-405B-Instruct","choices":[{"index":0,"text":" top destination for foodies, with","logprobs":null,"finish_reason":"length","stop_reason":null}],"usage":{"prompt_tokens":5,"total_tokens":12,"completion_tokens":7}}

Acelera los tiempos de carga de modelos con Hyperdisk ML (opcional)

El vLLM puede tardar hasta 90 minutos en descargar, cargar y activar Llama 3.1-405B en cada réplica nueva. Puedes reducir este tiempo a 20 minutos si descargas el modelo directamente en un Hyperdisk ML y lo conectas a cada pod.

Puedes seguir el instructivo Acelera la carga de datos de IA/AA con Hyperdisk ML con los siguientes archivos YAML:

  1. Guarda el siguiente manifiesto de ejemplo como producer-pvc.yaml:

    kind: PersistentVolumeClaim
    apiVersion: v1
    metadata:
      name: producer-pvc
    spec:
      storageClassName: hyperdisk-ml
      accessModes:
      - ReadWriteOnce
      resources:
        requests:
          storage: 800Gi
    
  2. Guarda el siguiente manifiesto de ejemplo como producer-job.yaml:

    apiVersion: batch/v1
    kind: Job
    metadata:
      name: producer-job
    spec:
      template:  # Template for the Pods the Job will create
        spec:
          affinity:
            nodeAffinity:
              requiredDuringSchedulingIgnoredDuringExecution:
                nodeSelectorTerms:
                - matchExpressions:
                  - key: cloud.google.com/machine-family
                    operator: In
                    values:
                    - "c3"
                - matchExpressions:
                  - key: topology.kubernetes.io/zone
                    operator: In
                    values:
                    - "ZONE"
          containers:
          - name: copy
            resources:
              requests:
                cpu: "32"
              limits:
                cpu: "32"
            image: python:3.11-bookworm
            command:
            - bash
            - -c
            - "pip install 'huggingface_hub==0.24.6' && \
              huggingface-cli download meta-llama/Meta-Llama-3.1-405B-Instruct --local-dir-use-symlinks=False --local-dir=/data/Meta-Llama-3.1-405B-Instruct --include *.safetensors *.json"
            env:
            - name: HUGGING_FACE_HUB_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            volumeMounts:
              - mountPath: "/data"
                name: volume
          restartPolicy: Never
          volumes:
            - name: volume
              persistentVolumeClaim:
                claimName: producer-pvc
      parallelism: 1         # Run 1 Pods concurrently
      completions: 1         # Once 1 Pods complete successfully, the Job is done
      backoffLimit: 4        # Max retries on failure
    

Implementa el servidor de modelos de vLLM

Después de completar los pasos, puedes implementar la implementación del servidor de GPU de varios nodos de vLLM que consume el volumen de Hyperdisk ML.



apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: vllm
spec:
  replicas: 1
  leaderWorkerTemplate:
    size: 2
    restartPolicy: RecreateGroupOnPodRestart
    leaderTemplate:
      metadata:
        labels:
          role: leader
      spec:
        containers:
          - name: vllm-leader
            image: IMAGE_NAME
            env:
              - name: HUGGING_FACE_HUB_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
            command:
              - sh
              - -c
              - "/workspace/vllm/examples/ray_init.sh leader --ray_cluster_size=$(LWS_GROUP_SIZE); 
                python3 -m vllm.entrypoints.openai.api_server --port 8080 --model /models/Meta-Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size 2"
            resources:
              limits:
                nvidia.com/gpu: "8"
            ports:
              - containerPort: 8080
            readinessProbe:
              tcpSocket:
                port: 8080
              initialDelaySeconds: 15
              periodSeconds: 10
            volumeMounts:
              - mountPath: /dev/shm
                name: dshm
              - mountPath: /models
                name: llama3-405b
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
        - name: llama3-405b
          persistentVolumeClaim:
            claimName: hdml-static-pvc
    workerTemplate:
      spec:
        containers:
          - name: vllm-worker
            image: IMAGE_NAME
            command:
              - sh
              - -c
              - "/workspace/vllm/examples/ray_init.sh worker --ray_address=$(LWS_LEADER_ADDRESS)"
            resources:
              limits:
                nvidia.com/gpu: "8"
            env:
              - name: HUGGING_FACE_HUB_TOKEN
                valueFrom:
                  secretKeyRef:
                    name: hf-secret
                    key: hf_api_token
            volumeMounts:
              - mountPath: /dev/shm
                name: dshm
              - mountPath: /models
                name: llama3-405b
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
        - name: llama3-405b
          persistentVolumeClaim:
            claimName: hdml-static-pvc
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-leader
spec:
  ports:
    - name: http
      port: 8080
      protocol: TCP
      targetPort: 8080
  selector:
    leaderworkerset.sigs.k8s.io/name: vllm
    role: leader
  type: ClusterIP

Limpia

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

Borra los recursos implementados

Para evitar que se generen cargos en tu cuenta de Google Cloud por los recursos que creaste en esta guía, ejecuta el siguiente comando:

gcloud container clusters delete CLUSTER_NAME \
  --location=ZONE

¿Qué sigue?