Ajustar modelos abiertos de Gemma con varias GPUs en GKE

En este tutorial se muestra cómo ajustar el modelo de lenguaje extenso (LLM) Gemma, una familia de modelos abiertos, usando unidades de procesamiento gráfico (GPUs) en Google Kubernetes Engine (GKE) con la biblioteca Transformers de Hugging Face. El ajuste fino es un proceso de aprendizaje supervisado que mejora la capacidad de un modelo preentrenado para realizar tareas específicas actualizando sus parámetros con un nuevo conjunto de datos. En este tutorial, descargarás los modelos preentrenados de la familia Gemma de 2B parámetros de Hugging Face y los ajustarás en un clúster Autopilot o Standard de GKE.

Esta guía es un buen punto de partida si necesitas el control granular, la escalabilidad, la resiliencia, la portabilidad y la rentabilidad de Kubernetes gestionado al ajustar un LLM.

Práctica recomendada:

Prueba nuestra solución Vertex AI si necesitas una plataforma de IA gestionada y unificada para crear y ofrecer modelos de aprendizaje automático rápidamente y de forma rentable.

Fondo

Si sirves Gemma mediante GPUs en GKE con la biblioteca de transformers, puedes implementar una solución de servicio de inferencia estable y lista para la 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.

Gemma

Gemma es un conjunto de modelos de IA generativa ligeros y disponibles públicamente que se han lanzado con una licencia abierta. Estos modelos de IA se pueden ejecutar en tus aplicaciones, hardware, dispositivos móviles o servicios alojados.

En esta guía, presentamos Gemma para la generación de texto. También puedes ajustar estos modelos para que se especialicen en realizar tareas específicas.

El conjunto de datos que se usa en este documento es b-mc2/sql-create-context.

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

GPUs

Las GPUs te permiten acelerar cargas de trabajo específicas que se ejecutan en tus nodos, como el aprendizaje automático y el procesamiento de datos. GKE ofrece una amplia gama de opciones de tipos de máquinas para la configuración de nodos, incluidos los tipos de máquinas con GPUs NVIDIA H100, L4 y A100.

Antes de usar GPUs en GKE, te recomendamos que completes el siguiente recorrido de aprendizaje:

  1. Consulta la disponibilidad de la versión actual de la GPU.
  2. Consulta información sobre las GPUs en GKE.

Hugging Face Transformers

Con la biblioteca Transformers de Hugging Face, puedes acceder a modelos preentrenados de vanguardia. La biblioteca Transformers te permite reducir el tiempo, los recursos y los costes computacionales asociados al entrenamiento completo de modelos.

En este tutorial, usarás las APIs y las herramientas de Hugging Face para descargar y perfeccionar estos modelos preentrenados.

Acceder al modelo

Para acceder a los modelos de Gemma e implementarlos en GKE, primero debes firmar el acuerdo de consentimiento de licencia y, a continuación, generar un token de acceso de Hugging Face.

Debes firmar el acuerdo de consentimiento para usar Gemma. Te las indicamos a continuación:

  1. Accede a la página de consentimiento del modelo en Kaggle.com.
  2. Verifica el consentimiento con tu cuenta de Hugging Face.
  3. Acepta los términos del modelo.

Generar un token de acceso

Para acceder al modelo a través de Hugging Face, necesitarás un token de Hugging Face.

Sigue estos pasos para generar un token si aún no tienes uno:

  1. Haz clic en Tu perfil > Configuración > Tokens de acceso.
  2. Selecciona New Token (Nuevo token).
  3. Especifica el nombre que quieras y un rol de al menos Write.
  4. Selecciona Generar un token.
  5. Copia el token generado en el portapapeles.

Prepara tu entorno

En este tutorial, usarás Cloud Shell para gestionar los recursos alojados enGoogle Cloud. Cloud Shell tiene preinstalado el software que necesitarás para este tutorial, como kubectl y la CLI de gcloud.

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

  1. En la Google Cloud consola, inicia una sesión de Cloud Shell haciendo clic en Icono de activación de Cloud Shell Activar Cloud Shell en la Google Cloud consola. Se iniciará una sesión en el panel inferior de la consola Google Cloud .

  2. Define las variables de entorno predeterminadas:

    gcloud config set project PROJECT_ID
    export PROJECT_ID=$(gcloud config get project)
    export CONTROL_PLANE_LOCATION=CONTROL_PLANE_LOCATION
    export CLUSTER_NAME=CLUSTER_NAME
    export HF_TOKEN=HF_TOKEN
    export HF_PROFILE=HF_PROFILE
    

    Sustituye los siguientes valores:

    • PROJECT_ID: tu Google Cloud ID de proyecto.
    • CONTROL_PLANE_LOCATION: la región de Compute Engine del plano de control de tu clúster. Proporciona una región que admita el tipo de acelerador que quieras usar. Por ejemplo, us-central1 para las GPUs L4.
    • CLUSTER_NAME: el nombre de tu clúster.
    • HF_TOKEN: el token de Hugging Face que has generado antes.
    • HF_PROFILE: el ID de perfil de Hugging Face que has creado anteriormente.
  3. Clona el repositorio de código de ejemplo de GitHub:

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
    cd kubernetes-engine-samples/ai-ml/llm-finetuning-gemma
    

Crear y configurar Google Cloud recursos

Sigue estas instrucciones para crear los recursos necesarios.

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

Puedes servir Gemma en GPUs en un clúster Autopilot o Standard de GKE. 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.

Práctica recomendada:

Usa Autopilot para disfrutar de una experiencia de Kubernetes totalmente gestionada.

Autopilot

En Cloud Shell, ejecuta el siguiente comando:

gcloud container clusters create-auto CLUSTER_NAME \
    --project=PROJECT_ID \
    --location=CONTROL_PLANE_LOCATION \
    --release-channel=rapid \
    --cluster-version=1.29

Sustituye los siguientes valores:

  • PROJECT_ID: tu Google Cloud ID de proyecto.
  • CONTROL_PLANE_LOCATION: la región de Compute Engine del plano de control de tu clúster. Proporciona una región que admita el tipo de acelerador que quieras usar. Por ejemplo, us-central1 para las GPUs L4.
  • CLUSTER_NAME: el nombre de tu clúster.

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

Estándar

  1. En Cloud Shell, ejecuta el siguiente comando para crear un clúster Standard:

    gcloud container clusters create CLUSTER_NAME \
        --project=PROJECT_ID \
        --location=CONTROL_PLANE_LOCATION \
        --workload-pool=PROJECT_ID.svc.id.goog \
        --release-channel=rapid \
        --num-nodes=1
    

    Sustituye los siguientes valores:

    • PROJECT_ID: tu Google Cloud ID de proyecto.
    • CONTROL_PLANE_LOCATION: la región de Compute Engine del plano de control de tu clúster. Proporciona una región que admita el tipo de acelerador que quieras usar. Por ejemplo, us-central1 para las GPUs L4.
    • CLUSTER_NAME: el nombre de tu clúster.

    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 gpupool \
        --accelerator type=nvidia-l4,count=8,gpu-driver-version=latest \
        --project=PROJECT_ID \
        --location=CONTROL_PLANE_LOCATION \
        --node-locations=CONTROL_PLANE_LOCATION-a \
        --cluster=CLUSTER_NAME \
        --machine-type=g2-standard-96 \
        --num-nodes=1
    

    GKE crea un único grupo de nodos que contiene dos GPUs L4 por nodo.

Crear un secreto de Kubernetes para las credenciales de Hugging Face

En Cloud Shell, haz lo siguiente:

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

    gcloud container clusters get-credentials CLUSTER_NAME \
        --location=CONTROL_PLANE_LOCATION
    

    Sustituye los siguientes valores:

    • CONTROL_PLANE_LOCATION: la región de Compute Engine del plano de control de tu clúster.
    • CLUSTER_NAME: el nombre de tu clúster.
  2. Crea un secreto 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 -
    

    Sustituye $HF_TOKEN por el token de Hugging Face que has generado antes o usa la variable de entorno si la has definido.

Crear un contenedor de ajuste fino con Docker y Cloud Build

Este contenedor usa el código de PyTorch y Hugging Face Transformers para optimizar el modelo de Gemma preentrenado.

  1. Crea un repositorio de Docker de Artifact Registry:

    gcloud artifacts repositories create gemma \
        --project=PROJECT_ID \
        --repository-format=docker \
        --location=us \
        --description="Gemma Repo"
    

    Sustituye PROJECT_ID por el ID de tu proyecto. Google Cloud

  2. Crea y envía la imagen:

    gcloud builds submit .
    
  3. Exporta el IMAGE_URL para usarlo más adelante en este tutorial.

    export IMAGE_URL=us-docker.pkg.dev/PROJECT_ID/gemma/finetune-gemma-gpu:1.0.0
    

Ejecutar una tarea de ajuste fino en GKE

En esta sección, desplegarás el trabajo de ajuste fino de Gemma. Un controlador de trabajo de Kubernetes crea uno o varios pods y se asegura de que ejecuten correctamente una tarea específica.

  1. Abre el archivo finetune.yaml.

    apiVersion: batch/v1
    kind: Job
    metadata:
      name: finetune-job
      namespace: default
    spec:
      backoffLimit: 2
      template:
        metadata:
          annotations:
            kubectl.kubernetes.io/default-container: finetuner
        spec:
          terminationGracePeriodSeconds: 600
          containers:
          - name: finetuner
            image: $IMAGE_URL
            resources:
              limits:
                nvidia.com/gpu: "8"
            env:
            - name: MODEL_NAME
              value: "google/gemma-2b"
            - name: NEW_MODEL
              value: "gemma-2b-sql-finetuned"
            - name: LORA_R
              value: "8"
            - name: LORA_ALPHA
              value: "16"
            - name: TRAIN_BATCH_SIZE
              value: "1"
            - name: EVAL_BATCH_SIZE
              value: "2"
            - name: GRADIENT_ACCUMULATION_STEPS
              value: "2"
            - name: DATASET_LIMIT
              value: "1000"
            - name: MAX_SEQ_LENGTH
              value: "512"
            - name: LOGGING_STEPS
              value: "5"
            - name: HF_TOKEN
              valueFrom:
                secretKeyRef:
                  name: hf-secret
                  key: hf_api_token
            volumeMounts:
            - mountPath: /dev/shm
              name: dshm
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
          nodeSelector:
            cloud.google.com/gke-accelerator: nvidia-l4
          restartPolicy: OnFailure
  2. Aplica el archivo de manifiesto para crear el trabajo de ajuste fino:

    envsubst < finetune.yaml | kubectl apply -f -
    

    Esta instrucción sustituye IMAGE_URL por la variable del manifiesto.

  3. Monitoriza el trabajo ejecutando el siguiente comando:

    watch kubectl get pods
    
  4. Consulta los registros del trabajo ejecutando el siguiente comando:

    kubectl logs job.batch/finetune-job -f
    

    El recurso Job descarga los datos del modelo y, a continuación, lo ajusta en las ocho GPUs. Este proceso puede tardar hasta 20 minutos.

  5. Una vez completado el trabajo, ve a tu cuenta de Hugging Face. En tu perfil de Hugging Face aparecerá un nuevo modelo llamado HF_PROFILE/gemma-2b-sql-finetuned.

Servir el modelo ajustado en GKE

En esta sección, desplegarás el contenedor vLLM para servir el modelo de Gemma. En este tutorial se usa un despliegue de Kubernetes para desplegar el contenedor vLLM. 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.

  1. Crea el siguiente archivo de manifiesto serve-gemma.yaml:

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: vllm-gemma-deployment
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: gemma-server
      template:
        metadata:
          labels:
            app: gemma-server
            ai.gke.io/model: gemma-2b
            ai.gke.io/inference-server: vllm
            examples.ai.gke.io/source: user-guide
        spec:
          containers:
          - name: inference-server
            image: docker.io/vllm/vllm-openai:v0.10.0
            resources:
              requests:
                cpu: "2"
                memory: "7Gi"
                ephemeral-storage: "10Gi"
                nvidia.com/gpu: 1
              limits:
                cpu: "2"
                memory: "7Gi"
                ephemeral-storage: "10Gi"
                nvidia.com/gpu: 1
            command: ["python3", "-m", "vllm.entrypoints.openai.api_server"]
            args:
            - --model=$(MODEL_ID)
            - --tensor-parallel-size=1
            env:
            - name: LD_LIBRARY_PATH
              value: ${LD_LIBRARY_PATH}:/usr/local/nvidia/lib64
            - name: MODEL_ID
              value: google/gemma-2b
            - 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
          nodeSelector:
            cloud.google.com/gke-accelerator: nvidia-l4
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: llm-service
    spec:
      selector:
        app: gemma-server
      type: ClusterIP
      ports:
        - protocol: TCP
          port: 8000
          targetPort: 8000
  2. Crea la variable de entorno para el nuevo MODEL_ID:

    export MODEL_ID=HF_PROFILE/gemma-2b-sql-finetuned
    

    Sustituye HF_PROFILE por el ID de perfil de Hugging Face que has creado anteriormente.

  3. Sustituye MODEL_ID en el archivo de manifiesto:

    sed -i "s|google/gemma-2b|$MODEL_ID|g" serve-gemma.yaml
    
  4. Aplica el archivo de manifiesto:

    kubectl apply -f serve-gemma.yaml
    

    Un pod del clúster descarga los pesos del modelo de Hugging Face e inicia el motor de servicio.

  5. Espera a que la implementación esté disponible:

    kubectl wait --for=condition=Available --timeout=700s deployment/vllm-gemma-deployment
    
  6. Consulta los registros de la implementación en ejecución:

    kubectl logs -f -l app=gemma-server
    

El recurso Deployment descarga los datos del modelo. Este proceso puede tardar unos minutos. El resultado debería ser similar al siguiente:

INFO 01-26 19:02:54 model_runner.py:689] Graph capturing finished in 4 secs.
INFO:     Started server process [1]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)

Asegúrate de que el modelo se haya descargado por completo antes de pasar a la siguiente sección.

Aplicar el modelo

En esta sección, interactúas con el modelo.

Configurar la redirección de puertos

Una vez que se haya implementado el modelo, ejecuta el siguiente comando para configurar el reenvío de puertos al modelo:

kubectl port-forward service/llm-service 8000:8000

El resultado debería ser similar al siguiente:

Forwarding from 127.0.0.1:8000 -> 8000

Interactuar con el modelo mediante curl

En una nueva sesión de terminal, usa curl para chatear con tu modelo:

El siguiente comando de ejemplo es para TGI:

USER_PROMPT="Question: What is the total number of attendees with age over 30 at kubecon eu? Context: CREATE TABLE attendees (name VARCHAR, age INTEGER, kubecon VARCHAR)"

curl -X POST http://localhost:8000/generate \
  -H "Content-Type: application/json" \
  -d @- <<EOF
{
    "prompt": "${USER_PROMPT}",
    "temperature": 0.1,
    "top_p": 1.0,
    "max_tokens": 24
}
EOF

En el siguiente resultado se muestra un ejemplo de la respuesta del modelo:

{"generated_text":" Answer: SELECT COUNT(age) FROM attendees WHERE age > 30 AND kubecon = 'eu'\n"}

En función de la consulta, puede que tengas que cambiar el max_token para obtener un mejor resultado. También puedes usar el modelo ajustado con instrucciones para disfrutar de una mejor experiencia de chat.