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 herramientas de Hugging Face para descargar y perfeccionar estos modelos preentrenados.

Objetivos

Esta guía está dirigida a usuarios nuevos o actuales de GKE, ingenieros de aprendizaje automático, ingenieros de MLOps (DevOps) o administradores de plataformas que quieran usar las funciones de orquestación de contenedores de Kubernetes para optimizar LLMs en hardware de GPU H100, A100 y L4.

Al final de esta guía, deberías poder seguir estos pasos:

  1. Prepara tu entorno con un clúster de GKE en modo Autopilot.
  2. Crea un contenedor de ajuste.
  3. Usa GPUs para optimizar el modelo Gemma 2B y subirlo a Hugging Face.

Antes de empezar

  • 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.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  • Verify that billing is enabled for your Google Cloud project.

  • Enable the required API.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    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 Conceder acceso.
    4. En el campo Nuevos principales, introduce tu identificador de usuario. Normalmente, se trata de la dirección de correo de una cuenta de Google.

    5. En la lista Selecciona un rol, elige un rol.
    6. Para conceder más roles, haz clic en Añadir otro rol y añade cada rol adicional.
    7. Haz clic en Guardar.
    8. 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, implementará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 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.

      Limpieza

      Para evitar que los recursos utilizados en este tutorial se cobren en tu cuenta de Google Cloud, elimina el proyecto que contiene los recursos o conserva el proyecto y elimina los recursos.

      Eliminar los recursos desplegados

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

      gcloud container clusters delete 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. 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.

      Siguientes pasos