Servir un LLM con varias GPUs en GKE

En este tutorial se muestra cómo desplegar y servir un modelo de lenguaje extenso (LLM) mediante varias GPUs en GKE para realizar inferencias eficientes y escalables. Crea un clúster de GKE que use varias GPUs de nivel 4 y prepara la infraestructura para servir cualquiera de los siguientes modelos:

El número de GPUs necesarias varía en función del formato de datos del modelo. En este tutorial, cada modelo usa dos GPUs L4. Para obtener más información, consulta Calcular la cantidad de GPUs.

Este tutorial está dirigido a ingenieros de aprendizaje automático, administradores y operadores de plataformas, y especialistas en datos e IA que quieran usar las funciones de orquestación de contenedores de Kubernetes para ofrecer LLMs. Para obtener más información sobre los roles habituales y las tareas de ejemplo a las que se hace referencia en el contenido, consulta Roles y tareas de usuario habituales de GKE. Google Cloud

Antes de leer esta página, asegúrese de que conoce los siguientes conceptos:

Prepara tu entorno

  1. En la Google Cloud consola, inicia una instancia de Cloud Shell:
    Abrir Cloud Shell

  2. Define las variables de entorno predeterminadas:

    gcloud config set project PROJECT_ID
gcloud config set billing/quota_project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)
export CONTROL_PLANE_LOCATION=us-central1

    Sustituye PROJECT_ID por el Google Cloud ID de tu proyecto.

Crear un clúster y un grupo de nodos de GKE

Puedes servir LLMs en GPUs 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

  1. En Cloud Shell, ejecuta el siguiente comando:

    gcloud container clusters create-auto l4-demo \
  --project=${PROJECT_ID} \
  --location=${CONTROL_PLANE_LOCATION} \
  --release-channel=rapid

    GKE crea un clúster de Autopilot con nodos de CPU y GPU según lo soliciten las cargas de trabajo desplegadas.

  2. Configura kubectl para que se comunique con tu clúster:

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

Estándar

  1. En Cloud Shell, ejecuta el siguiente comando para crear un clúster estándar que utilice Workload Identity Federation for GKE:

    gcloud container clusters create l4-demo \
  --location ${CONTROL_PLANE_LOCATION} \
  --workload-pool ${PROJECT_ID}.svc.id.goog \
  --enable-image-streaming \
  --node-locations=${CONTROL_PLANE_LOCATION}-a \
  --workload-pool=${PROJECT_ID}.svc.id.goog \
  --machine-type n2d-standard-4 \
  --num-nodes 1 --min-nodes 1 --max-nodes 5 \
  --release-channel=rapid

    La creación del clúster puede tardar varios minutos.

  2. Ejecuta el siguiente comando para crear un grupo de nodos para tu clúster:

    gcloud container node-pools create g2-standard-24 --cluster l4-demo \
  --location ${CONTROL_PLANE_LOCATION} \
  --accelerator type=nvidia-l4,count=2,gpu-driver-version=latest \
  --machine-type g2-standard-24 \
  --enable-autoscaling --enable-image-streaming \
  --num-nodes=0 --min-nodes=0 --max-nodes=3 \
  --node-locations ${CONTROL_PLANE_LOCATION}-a,${CONTROL_PLANE_LOCATION}-c \
  --spot

    GKE crea los siguientes recursos para el LLM:

    • Un clúster público de la edición Estándar de Google Kubernetes Engine (GKE).
    • Un grupo de nodos con el tipo de máquina g2-standard-24 se ha reducido a 0 nodos. No se te cobra por ninguna GPU hasta que inicies pods que soliciten GPUs. Este grupo de nodos aprovisiona máquinas virtuales de acceso puntual, que tienen un precio más bajo que las máquinas virtuales estándar de Compute Engine predeterminadas y no ofrecen ninguna garantía de disponibilidad. Puedes quitar la marca --spot de este comando y el selector de nodos cloud.google.com/gke-spot de la configuración text-generation-inference.yaml para usar las VMs bajo demanda.

  3. Configura kubectl para que se comunique con tu clúster:

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

Preparar la carga de trabajo

En esta sección se explica cómo configurar tu carga de trabajo en función del modelo que quieras usar. En este tutorial se usan implementaciones de Kubernetes para desplegar el modelo. 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.

Llama 3 70b

  1. Define las variables de entorno predeterminadas:

    export HF_TOKEN=HUGGING_FACE_TOKEN

    Sustituye HUGGING_FACE_TOKEN por tu token de Hugging Face.

  2. Crea un secreto de Kubernetes para el token de Hugging Face:

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

  3. Crea el siguiente archivo de manifiesto de text-generation-inference.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: llm
spec:
  replicas: 1
  selector:
    matchLabels:
      app: llm
  template:
    metadata:
      labels:
        app: llm
    spec:
      containers:
      - name: llm
        image: us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310
        resources:
          requests:
            cpu: "10"
            memory: "60Gi"
            nvidia.com/gpu: "2"
          limits:
            cpu: "10"
            memory: "60Gi"
            nvidia.com/gpu: "2"
        env:
        - name: MODEL_ID
          value: meta-llama/Meta-Llama-3-70B-Instruct
        - name: NUM_SHARD
          value: "2"
        - name: MAX_INPUT_TOKENS
          value: "2048"
        - name: PORT
          value: "8080"
        - name: QUANTIZE
          value: bitsandbytes-nf4
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              name: l4-demo
              key: HUGGING_FACE_TOKEN
        volumeMounts:
          - mountPath: /dev/shm
            name: dshm
          # mountPath is set to /tmp as it's the path where the HUGGINGFACE_HUB_CACHE environment
          # variable in the TGI DLCs is set to instead of the default /data set within the TGI default image.
          # i.e. where the downloaded model from the Hub will be stored
          - mountPath: /tmp
            name: ephemeral-volume
      volumes:
        - name: dshm
          emptyDir:
              medium: Memory
        - name: ephemeral-volume
          ephemeral:
            volumeClaimTemplate:
              metadata:
                labels:
                  type: ephemeral
              spec:
                accessModes: ["ReadWriteOnce"]
                storageClassName: "premium-rwo"
                resources:
                  requests:
                    storage: 150Gi
      nodeSelector:
        cloud.google.com/gke-accelerator: "nvidia-l4"
        cloud.google.com/gke-spot: "true"

    En este manifiesto:

    • NUM_SHARD debe ser 2 porque el modelo requiere dos GPUs NVIDIA L4.
    • QUANTIZE se ha definido como bitsandbytes-nf4, lo que significa que el modelo se ha cargado en 4 bits en lugar de 32. De esta forma, GKE puede reducir la cantidad de memoria de GPU necesaria y mejorar la velocidad de inferencia. Sin embargo, la precisión del modelo puede disminuir. Para saber cómo calcular las GPUs que debes solicitar, consulta Calcular la cantidad de GPUs.

  4. Aplica el archivo de manifiesto:

    kubectl apply -f text-generation-inference.yaml

    El resultado debería ser similar al siguiente:

    deployment.apps/llm created

  5. Verifica el estado del modelo:

    kubectl get deploy

    El resultado debería ser similar al siguiente:

    NAME          READY   UP-TO-DATE   AVAILABLE   AGE
llm           1/1     1            1           20m

  6. Consulta los registros de la implementación en ejecución:

    kubectl logs -l app=llm

    El resultado debería ser similar al siguiente:

    {"timestamp":"2024-03-09T05:08:14.751646Z","level":"INFO","message":"Warming up model","target":"text_generation_router","filename":"router/src/main.rs","line_number":291}
{"timestamp":"2024-03-09T05:08:19.961136Z","level":"INFO","message":"Setting max batch total tokens to 133696","target":"text_generation_router","filename":"router/src/main.rs","line_number":328}
{"timestamp":"2024-03-09T05:08:19.961164Z","level":"INFO","message":"Connected","target":"text_generation_router","filename":"router/src/main.rs","line_number":329}
{"timestamp":"2024-03-09T05:08:19.961171Z","level":"WARN","message":"Invalid hostname, defaulting to 0.0.0.0","target":"text_generation_router","filename":"router/src/main.rs","line_number":343}

Mixtral 8x7b

  1. Define las variables de entorno predeterminadas:

    export HF_TOKEN=HUGGING_FACE_TOKEN

    Sustituye HUGGING_FACE_TOKEN por tu token de Hugging Face.

  2. Crea un secreto de Kubernetes para el token de Hugging Face:

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

  3. Crea el siguiente archivo de manifiesto de text-generation-inference.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: llm
spec:
  replicas: 1
  selector:
    matchLabels:
      app: llm
  template:
    metadata:
      labels:
        app: llm
    spec:
      containers:
      - name: llm
        image: us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-3.ubuntu2204.py311
        resources:
          requests:
            cpu: "5"
            memory: "40Gi"
            nvidia.com/gpu: "2"
          limits:
            cpu: "5"
            memory: "40Gi"
            nvidia.com/gpu: "2"
        env:
        - name: MODEL_ID
          value: mistralai/Mixtral-8x7B-Instruct-v0.1
        - name: NUM_SHARD
          value: "2"
        - name: PORT
          value: "8080"
        - name: QUANTIZE
          value: bitsandbytes-nf4
        - name: HUGGING_FACE_HUB_TOKEN
          valueFrom:
            secretKeyRef:
              name: l4-demo
              key: HUGGING_FACE_TOKEN          
        volumeMounts:
          - mountPath: /dev/shm
            name: dshm
          # mountPath is set to /tmp as it's the path where the HF_HOME environment
          # variable in the TGI DLCs is set to instead of the default /data set within the TGI default image.
          # i.e. where the downloaded model from the Hub will be stored
          - mountPath: /tmp
            name: ephemeral-volume
      volumes:
        - name: dshm
          emptyDir:
              medium: Memory
        - name: ephemeral-volume
          ephemeral:
            volumeClaimTemplate:
              metadata:
                labels:
                  type: ephemeral
              spec:
                accessModes: ["ReadWriteOnce"]
                storageClassName: "premium-rwo"
                resources:
                  requests:
                    storage: 100Gi
      nodeSelector:
        cloud.google.com/gke-accelerator: "nvidia-l4"
        cloud.google.com/gke-spot: "true"

    En este manifiesto:

    • NUM_SHARD debe ser 2 porque el modelo requiere dos GPUs NVIDIA L4.
    • QUANTIZE se ha definido como bitsandbytes-nf4, lo que significa que el modelo se ha cargado en 4 bits en lugar de 32. De esta forma, GKE puede reducir la cantidad de memoria de GPU necesaria y mejorar la velocidad de inferencia. Sin embargo, esto puede reducir la precisión del modelo. Para saber cómo calcular las GPUs que debes solicitar, consulta Calcular la cantidad de GPUs.

  4. Aplica el archivo de manifiesto:

    kubectl apply -f text-generation-inference.yaml

    El resultado debería ser similar al siguiente:

    deployment.apps/llm created

  5. Verifica el estado del modelo:

    watch kubectl get deploy

    Cuando la implementación esté lista, el resultado será similar al siguiente:

    NAME          READY   UP-TO-DATE   AVAILABLE   AGE
llm           1/1     1            1           10m

    Para salir del reloj, escribe CTRL + C.

  6. Consulta los registros de la implementación en ejecución:

    kubectl logs -l app=llm

    El resultado debería ser similar al siguiente:

    {"timestamp":"2024-03-09T05:08:14.751646Z","level":"INFO","message":"Warming up model","target":"text_generation_router","filename":"router/src/main.rs","line_number":291}
{"timestamp":"2024-03-09T05:08:19.961136Z","level":"INFO","message":"Setting max batch total tokens to 133696","target":"text_generation_router","filename":"router/src/main.rs","line_number":328}
{"timestamp":"2024-03-09T05:08:19.961164Z","level":"INFO","message":"Connected","target":"text_generation_router","filename":"router/src/main.rs","line_number":329}
{"timestamp":"2024-03-09T05:08:19.961171Z","level":"WARN","message":"Invalid hostname, defaulting to 0.0.0.0","target":"text_generation_router","filename":"router/src/main.rs","line_number":343}

Falcon 40b

  1. Crea el siguiente archivo de manifiesto de text-generation-inference.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: llm
spec:
  replicas: 1
  selector:
    matchLabels:
      app: llm
  template:
    metadata:
      labels:
        app: llm
    spec:
      containers:
      - name: llm
        image: us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.1-4.ubuntu2204.py310
        resources:
          requests:
            cpu: "10"
            memory: "60Gi"
            nvidia.com/gpu: "2"
          limits:
            cpu: "10"
            memory: "60Gi"
            nvidia.com/gpu: "2"
        env:
        - name: MODEL_ID
          value: tiiuae/falcon-40b-instruct
        - name: NUM_SHARD
          value: "2"
        - name: PORT
          value: "8080"
        - name: QUANTIZE
          value: bitsandbytes-nf4
        volumeMounts:
          - mountPath: /dev/shm
            name: dshm
          # mountPath is set to /data as it's the path where the HUGGINGFACE_HUB_CACHE environment
          # variable points to in the TGI container image i.e. where the downloaded model from the Hub will be
          # stored
          - mountPath: /data
            name: ephemeral-volume
      volumes:
        - name: dshm
          emptyDir:
              medium: Memory
        - name: ephemeral-volume
          ephemeral:
            volumeClaimTemplate:
              metadata:
                labels:
                  type: ephemeral
              spec:
                accessModes: ["ReadWriteOnce"]
                storageClassName: "premium-rwo"
                resources:
                  requests:
                    storage: 175Gi
      nodeSelector:
        cloud.google.com/gke-accelerator: "nvidia-l4"
        cloud.google.com/gke-spot: "true"

    En este manifiesto:

    • NUM_SHARD debe ser 2 porque el modelo requiere dos GPUs NVIDIA L4.
    • QUANTIZE se ha definido como bitsandbytes-nf4, lo que significa que el modelo se ha cargado en 4 bits en lugar de 32. De esta forma, GKE puede reducir la cantidad de memoria de GPU necesaria y mejorar la velocidad de inferencia. Sin embargo, la precisión del modelo puede disminuir. Para saber cómo calcular las GPUs que debes solicitar, consulta Calcular la cantidad de GPUs.

  2. Aplica el archivo de manifiesto:

    kubectl apply -f text-generation-inference.yaml

    El resultado debería ser similar al siguiente:

    deployment.apps/llm created

  3. Verifica el estado del modelo:

    watch kubectl get deploy

    Cuando la implementación esté lista, el resultado será similar al siguiente:

    NAME          READY   UP-TO-DATE   AVAILABLE   AGE
llm           1/1     1            1           10m

    Para salir del reloj, escribe CTRL + C.

  4. Consulta los registros de la implementación en ejecución:

    kubectl logs -l app=llm

    El resultado debería ser similar al siguiente:

    {"timestamp":"2024-03-09T05:08:14.751646Z","level":"INFO","message":"Warming up model","target":"text_generation_router","filename":"router/src/main.rs","line_number":291}
{"timestamp":"2024-03-09T05:08:19.961136Z","level":"INFO","message":"Setting max batch total tokens to 133696","target":"text_generation_router","filename":"router/src/main.rs","line_number":328}
{"timestamp":"2024-03-09T05:08:19.961164Z","level":"INFO","message":"Connected","target":"text_generation_router","filename":"router/src/main.rs","line_number":329}
{"timestamp":"2024-03-09T05:08:19.961171Z","level":"WARN","message":"Invalid hostname, defaulting to 0.0.0.0","target":"text_generation_router","filename":"router/src/main.rs","line_number":343}

Crear un servicio de tipo ClusterIP

Expón tus pods internamente en el clúster para que otras aplicaciones puedan descubrirlos y acceder a ellos.

  1. Crea el siguiente archivo de manifiesto llm-service.yaml:

    apiVersion: v1
kind: Service
metadata:
  name: llm-service
spec:
  selector:
    app: llm
  type: ClusterIP
  ports:
    - protocol: TCP
      port: 80
      targetPort: 8080

  2. Aplica el archivo de manifiesto:

    kubectl apply -f llm-service.yaml

Implementar una interfaz de chat

Usa Gradio para crear una aplicación web que te permita interactuar con tu modelo. Gradio es una biblioteca de Python que tiene un wrapper ChatInterface que crea interfaces de usuario para chatbots.

Llama 3 70b

  1. Crea un archivo llamado gradio.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: gradio
  labels:
    app: gradio
spec:
  strategy:
    type: Recreate
  replicas: 1
  selector:
    matchLabels:
      app: gradio
  template:
    metadata:
      labels:
        app: gradio
    spec:
      containers:
      - name: gradio
        image: us-docker.pkg.dev/google-samples/containers/gke/gradio-app:v1.0.4
        resources:
          requests:
            cpu: "512m"
            memory: "512Mi"
          limits:
            cpu: "1"
            memory: "512Mi"
        env:
        - name: CONTEXT_PATH
          value: "/generate"
        - name: HOST
          value: "http://llm-service"
        - name: LLM_ENGINE
          value: "tgi"
        - name: MODEL_ID
          value: "meta-llama/Meta-Llama-3-70B-Instruct"
        - name: USER_PROMPT
          value: "<|begin_of_text|><|start_header_id|>user<|end_header_id|> prompt <|eot_id|><|start_header_id|>assistant<|end_header_id|>"
        - name: SYSTEM_PROMPT
          value: "prompt <|eot_id|>"
        ports:
        - containerPort: 7860
---
apiVersion: v1
kind: Service
metadata:
  name: gradio-service
spec:
  type: LoadBalancer
  selector:
    app: gradio
  ports:
  - port: 80
    targetPort: 7860

  2. Aplica el archivo de manifiesto:

    kubectl apply -f gradio.yaml

  3. Busca la dirección IP externa del servicio:

    kubectl get svc

    El resultado debería ser similar al siguiente:

    NAME             TYPE           CLUSTER-IP     EXTERNAL-IP     PORT(S)        AGE
gradio-service   LoadBalancer   10.24.29.197   34.172.115.35   80:30952/TCP   125m

  4. Copia la dirección IP externa de la columna EXTERNAL-IP.

  5. Para ver la interfaz del modelo en tu navegador web, usa la dirección IP externa con el puerto expuesto:

    http://EXTERNAL_IP

Mixtral 8x7b

  1. Crea un archivo llamado gradio.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: gradio
  labels:
    app: gradio
spec:
  strategy:
    type: Recreate
  replicas: 1
  selector:
    matchLabels:
      app: gradio
  template:
    metadata:
      labels:
        app: gradio
    spec:
      containers:
      - name: gradio
        image: us-docker.pkg.dev/google-samples/containers/gke/gradio-app:v1.0.4
        resources:
          requests:
            cpu: "512m"
            memory: "512Mi"
          limits:
            cpu: "1"
            memory: "512Mi"
        env:
        - name: CONTEXT_PATH
          value: "/generate"
        - name: HOST
          value: "http://llm-service"
        - name: LLM_ENGINE
          value: "tgi"
        - name: MODEL_ID
          value: "mixtral-8x7b"
        - name: USER_PROMPT
          value: "[INST] prompt [/INST]"
        - name: SYSTEM_PROMPT
          value: "prompt"
        ports:
        - containerPort: 7860
---
apiVersion: v1
kind: Service
metadata:
  name: gradio-service
spec:
  type: LoadBalancer
  selector:
    app: gradio
  ports:
  - port: 80
    targetPort: 7860

  2. Aplica el archivo de manifiesto:

    kubectl apply -f gradio.yaml

  3. Busca la dirección IP externa del servicio:

    kubectl get svc

    El resultado debería ser similar al siguiente:

    NAME             TYPE           CLUSTER-IP     EXTERNAL-IP     PORT(S)        AGE
gradio-service   LoadBalancer   10.24.29.197   34.172.115.35   80:30952/TCP   125m

  4. Copia la dirección IP externa de la columna EXTERNAL-IP.

  5. Para ver la interfaz del modelo en tu navegador web, usa la dirección IP externa con el puerto expuesto:

    http://EXTERNAL_IP

Falcon 40b

  1. Crea un archivo llamado gradio.yaml:

    apiVersion: apps/v1
kind: Deployment
metadata:
  name: gradio
  labels:
    app: gradio
spec:
  strategy:
    type: Recreate
  replicas: 1
  selector:
    matchLabels:
      app: gradio
  template:
    metadata:
      labels:
        app: gradio
    spec:
      containers:
      - name: gradio
        image: us-docker.pkg.dev/google-samples/containers/gke/gradio-app:v1.0.4
        resources:
          requests:
            cpu: "512m"
            memory: "512Mi"
          limits:
            cpu: "1"
            memory: "512Mi"
        env:
        - name: CONTEXT_PATH
          value: "/generate"
        - name: HOST
          value: "http://llm-service"
        - name: LLM_ENGINE
          value: "tgi"
        - name: MODEL_ID
          value: "falcon-40b-instruct"
        - name: USER_PROMPT
          value: "User: prompt"
        - name: SYSTEM_PROMPT
          value: "Assistant: prompt"
        ports:
        - containerPort: 7860
---
apiVersion: v1
kind: Service
metadata:
  name: gradio-service
spec:
  type: LoadBalancer
  selector:
    app: gradio
  ports:
  - port: 80
    targetPort: 7860

  2. Aplica el archivo de manifiesto:

    kubectl apply -f gradio.yaml

  3. Busca la dirección IP externa del servicio:

    kubectl get svc

    El resultado debería ser similar al siguiente:

    NAME             TYPE           CLUSTER-IP     EXTERNAL-IP     PORT(S)        AGE
gradio-service   LoadBalancer   10.24.29.197   34.172.115.35   80:30952/TCP   125m

  4. Copia la dirección IP externa de la columna EXTERNAL-IP.

  5. Para ver la interfaz del modelo en tu navegador web, usa la dirección IP externa con el puerto expuesto:

    http://EXTERNAL_IP

Calcular la cantidad de GPUs

La cantidad de GPUs depende del valor de la marca QUANTIZE. En este tutorial, QUANTIZE se ha definido como bitsandbytes-nf4, lo que significa que el modelo se carga en 4 bits.

Un modelo de 70.000 millones de parámetros requeriría un mínimo de 40 GB de memoria de GPU, lo que equivale a 70.000 millones multiplicado por 4 bits (70.000 millones × 4 bits= 35 GB) y considera una sobrecarga de 5 GB. En este caso, una sola GPU L4 no tendría suficiente memoria. Por lo tanto, en los ejemplos de este tutorial se usan dos GPUs L4 de memoria (2 x 24 = 48 GB). Esta configuración es suficiente para ejecutar Falcon 40b o Llama 3 70b en GPUs L4.