Entrega modelos abiertos de Gemma con TPU en GKE con Saxml


En esta guía, se muestra cómo entregar un modelo de lenguaje grande (LLM) de modelos abiertos de Gemma con unidades de procesamiento tensorial (TPU) en Google Kubernetes Engine (GKE) con Saxml. En esta guía, debes descargar los modelos de Gemma ajustados con instrucciones de parámetros 2B y 7B en Cloud Storage y, luego, implementarlos en un clúster de GKE Standard mediante contenedores que ejecutan Saxml.

Esta guía es un buen punto de partida si necesitas la escalabilidad, la resiliencia y la rentabilidad que ofrecen las funciones de Kubernetes cuando implementas tu modelo en Saxml.

Si necesitas una plataforma de IA administrada unificada para compilar y entregar modelos de AA con rapidez de forma rentable, te recomendamos que pruebes nuestra solución de implementación de Vertex AI.

Antecedentes

La entrega de Gemma mediante TPU en GKE con Saxml permite implementar una solución de entrega de inferencia sólida y lista para la producción con todos los beneficios de administrarKubernetes, incluida la escalabilidad eficiente y la mayor disponibilidad. En esta sección, se describen las tecnologías clave que se usan en este instructivo.

Gemma

Gemma es un conjunto de modelos de IA generativa básicos y de disponibilidad general que se lanzan con una licencia abierta. Estos modelos de IA están disponibles para ejecutarse en tus aplicaciones, hardware, dispositivos móviles o servicios alojados. Puedes usar los modelos de Gemma para la generación de texto, y puedes ajustar estos modelos en el caso de tareas especializadas.

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

TPU

Las TPUs son circuitos integrados personalizados específicos de aplicaciones (ASIC) de Google que se usan para acelerar los frameworks de procesamiento de datos como TensorFlow, PyTorch y JAX.

Antes de usar las TPU en GKE, te recomendamos que completes la siguiente ruta de aprendizaje:

  1. Obtén información sobre la disponibilidad actual de la versión de TPU con la arquitectura del sistema de Cloud TPU.
  2. Obtén información sobre las TPU en GKE.

En este instructivo, se implementan los modelos Gemma 2B y Gemma 7B. GKE aloja estos modelos en los siguientes grupos de nodos de TPU v5e de host único:

  • Gemma 2B: Modelo ajustado de instrucciones alojado en un grupo de nodos TPU v5e con una topología 1x1 que representa un chip TPU. El tipo de máquina para los nodos es ct5lp-hightpu-1t.
  • Gemma 7B: Modelo ajustado de instrucciones alojado en un grupo de nodos TPU v5e con una topología 2x2 que representa cuatro chips TPU. El tipo de máquina para los nodos es ct5lp-hightpu-4t.

Saxml

Saxml es un sistema experimental que admite modelos Paxml, JAX y PyTorch para la inferencia. El sistema Saxml incluye los siguientes componentes:

  • Celda de Sax o clúster de Sax: Un servidor de administración y un grupo de servidores de modelos. El servidor de administración realiza un seguimiento de los servidores de modelos, asigna modelos publicados a los servidores de modelos para entregarlos y ayuda a los clientes a ubicar los servidores de modelos que entregan modelos publicados específicos.
  • Cliente de Saxml: La interfaz de programación para el usuario para el sistema Saxml. El cliente de Saxml incluye una herramienta de línea de comandos (saxutil) y un conjunto de bibliotecas cliente en Python, C++ y Go.

En este instructivo, también usarás el servidor HTTP de Saxml. El servidor HTTP de Saxml es un servidor HTTP personalizado que encapsula la biblioteca cliente de Python de Saxml y expone a las APIs de REST para interactuar con el sistema Saxml. Las APIs de REST incluyen extremos para publicar, enumerar, anular la publicación de modelos y generar predicciones.

Objetivos

Este instructivo está dirigido a clientes de IA generativa que usan JAX, además de usuarios nuevos o existentes de GKE que deseen usar las funciones de organización de contenedores de Kubernetes para entregar Gemma, como ingenieros de AA, MLOps (DevOps).

En este instructivo, se abarcan los siguientes pasos:

  1. Prepara un clúster de GKE Standard con la topología de TPU recomendada según las características del modelo.
  2. Implementa componentes de Saxml en GKE.
  3. Obtén y publica el modelo de parámetros Gemma 2B o Gemma 7B.
  4. Entrega los modelos publicados e interactuar con ellos.

Arquitectura

En esta sección, se describe la arquitectura de GKE que se usa en este instructivo. La arquitectura consta de un clúster de GKE Standard que aprovisiona TPU y aloja componentes de Saxml para implementar y entregar modelos Gemma 2B o 7B. En el siguiente diagrama, se muestran los componentes de esta arquitectura:

Un diagrama de la arquitectura implementada en este instructivo

Esta arquitectura incluye los siguientes componentes:

  • Un clúster zonal de GKE Standard.
  • Un grupo de nodos de porción de TPU de host único que depende del modelo de Gemma que deseas entregar:
    • Gemma 2B: Se configura con una TPU v5e con una topología 1x1. Se configuró una instancia del servidor de modelos de Saxml para usar este grupo de nodos.
    • Gemma 7B: Configurado con una TPU v5e con una topología 2x2. Se configuró una instancia del servidor de modelos de Saxml para usar este grupo de nodos.
  • Un grupo de nodos de CPU predeterminado en el que se implementan el servidor de administración de Saxml y el servidor HTTP de Saxml.
  • Dos buckets de Cloud Storage:
    • Un bucket de Cloud Storage almacena el estado que administra un servidor de administración.
    • Un bucket de Cloud Storage almacena los puntos de control del modelo de Gemma.

Esta arquitectura tiene las siguientes características:

  • Un Artifact Registry público administra las imágenes de contenedor para los componentes de Saxml.
  • El clúster de GKE usa la federación de identidades para cargas de trabajo para GKE. Todos los componentes de Saxml usan una federación de identidades para cargas de trabajo que integra una cuenta de servicio de IAM para acceder a servicios externos, como los buckets de Cloud Storage.
  • Los registros que generan los componentes de Saxml están integrados en Cloud Logging.
  • Puedes usar Cloud Monitoring para analizar las métricas de rendimiento de los grupos de nodos de GKE, como se crea en este instructivo.

Antes de comenzar

  • Accede a tu cuenta de Google Cloud. Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.
  • In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  • Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.

  • 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

  • Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.

  • 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 colunn 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.
  • Asegúrate de tener suficientes cuotas para 5 chips TPU v5e. En este instructivo, usarás instancias bajo demanda.
  • Crea una cuenta de Kaggle, si todavía no tienes una.

Prepara el entorno para Gemma

Inicia Cloud Shell

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

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

  2. Configura las variables de entorno predeterminadas:

    gcloud config set project PROJECT_ID
    export PROJECT_ID=$(gcloud config get project)
    export LOCATION=LOCATION
    export CLUSTER_NAME=saxml-tpu
    

    Reemplaza los siguientes valores:

Crea un clúster de GKE Standard

En esta sección, crearás el clúster de GKE y el grupo de nodos.

Gemma 2B-it

Usa Cloud Shell para realizar las siguientes acciones:

  1. Crea un clúster estándar que use la federación de identidades para cargas de trabajo para GKE:

    gcloud container clusters create ${CLUSTER_NAME} \
        --enable-ip-alias \
        --machine-type=e2-standard-4 \
        --num-nodes=2 \
        --release-channel=rapid \
        --workload-pool=${PROJECT_ID}.svc.id.goog \
        --location=${LOCATION}
    

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

  2. Crea un grupo de nodos TPU v5e con una topología 1x1 y un nodo:

    gcloud container node-pools create tpu-v5e-1x1 \
        --cluster=${CLUSTER_NAME} \
        --machine-type=ct5lp-hightpu-1t \
        --num-nodes=1 \
        --location=${LOCATION}
    

    Debes entregar el modelo de Gemma 2B en este grupo de nodos.

Gemma 7B-it

Usa Cloud Shell para realizar las siguientes acciones:

  1. Crea un clúster estándar que use la federación de identidades para cargas de trabajo para GKE:

    gcloud container clusters create ${CLUSTER_NAME} \
        --enable-ip-alias \
        --machine-type=e2-standard-4 \
        --num-nodes=2 \
        --release-channel=rapid \
        --workload-pool=${PROJECT_ID}.svc.id.goog \
        --location=${LOCATION}
    

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

  2. Crea un grupo de nodos TPU v5e con una topología 2x2 y un nodo:

    gcloud container node-pools create tpu-v5e-2x2 \
        --cluster=${CLUSTER_NAME} \
        --machine-type=ct5lp-hightpu-4t \
        --num-nodes=1 \
        --location=${LOCATION}
    

    Debes entregar el modelo de Gemma 7B en este grupo de nodos.

Crea los buckets de Cloud Storage

Crea dos buckets de Cloud Storage para administrar el estado del servidor de administración de Saxml y los puntos de control del modelo.

En Cloud Shell, ejecuta lo siguiente:

  1. Crea un bucket de Cloud Storage para almacenar los parámetros de configuración del servidor de administración de Saxml.

    gcloud storage buckets create gs://ADMIN_BUCKET_NAME
    

    Reemplaza ADMIN_BUCKET_NAME por el nombre del bucket de Cloud Storage que almacena el servidor de administración de Saxml.

  2. Crea un bucket de Cloud Storage para almacenar puntos de control de modelos:

    gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME
    

    Reemplaza CHECKPOINTS_BUCKET_NAME por el nombre del Bucket de Cloud Storage que almacena los puntos de control del modelo.

Configura el acceso a tus cargas de trabajo mediante la federación de identidades para cargas de trabajo para GKE

Asigna una Cuenta de servicio de Kubernetes a la aplicación y configúrala para que actúe como una cuenta de servicio de IAM.

  1. Configura kubectl para comunicarse con tu clúster:

    gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}
    
  2. Crea una cuenta de servicio de Kubernetes para que tu aplicación use:

    gcloud iam service-accounts create wi-sax
    
  3. Agrega una vinculación de política de IAM para que tu cuenta de servicio de IAM lea y escriba en Cloud Storage:

    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
        --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
        --role roles/storage.objectUser
    
    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
        --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
        --role roles/storage.insightsCollectorService
    
  4. Para permitir que la cuenta de servicio de Kubernetes actúe en nombre de la cuenta de servicio de IAM, agrega una vinculación de política de IAM entre las dos. Esta vinculación permite que la cuenta de servicio de Kubernetes actúe como la cuenta de servicio de IAM.

    gcloud iam service-accounts add-iam-policy-binding wi-sax@${PROJECT_ID}.iam.gserviceaccount.com \
        --role roles/iam.workloadIdentityUser \
        --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"
    
  5. Anota la cuenta de servicio de Kubernetes con la dirección de correo electrónico de la cuenta de servicio de IAM.

    kubectl annotate serviceaccount default \
        iam.gke.io/gcp-service-account=wi-sax@${PROJECT_ID}.iam.gserviceaccount.com
    

Obtén acceso al modelo

Para obtener acceso a los modelos de Gemma para la implementación en GKE, debes acceder a la plataforma de Kaggle, firmar el contrato de consentimiento de licencia y obtener un token de la API de Kaggle. En este instructivo, usas un Secret de Kubernetes para las credenciales de Kaggle.

Debes firmar el acuerdo de consentimiento para usar Gemma. Sigue estas instrucciones:

  1. Accede a la página de consentimiento del modelo en Kaggle.com.
  2. Ingresa en Kaggle si aún no lo has hecho.
  3. Haz clic en Solicitar acceso.
  4. En la sección Elegir cuenta para el consentimiento, selecciona Verificar mediante la cuenta de Kaggle para usar tu cuenta de Kaggle para obtener el consentimiento.
  5. Acepta los Términos y Condiciones del modelo.

Genera un token de acceso

Para acceder al modelo a través de Kaggle, necesitas un token de la API de Kaggle.

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

  1. En el navegador, ve a la configuración de Kaggle.
  2. En la sección API, haz clic en Crear token nuevo.

Kaggle descarga un archivo llamado kaggle.json.

Sube el token de acceso a Cloud Shell

En Cloud Shell, puedes subir el token de la API de Kaggle a tu proyecto de Google Cloud:

  1. En Cloud Shell, haz clic en Más > Subir.
  2. Elige Archivo y haz clic en Elegir archivos.
  3. Abre el archivo kaggle.json.
  4. Haz clic en Subir.

Crea un secreto de Kubernetes para las credenciales de Kaggle

En Cloud Shell, haz lo siguiente:

  1. Configura kubectl para comunicarse con tu clúster:

    gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}
    
  2. Crea un secreto para almacenar las credenciales de Kaggle:

    kubectl create secret generic kaggle-secret \
        --from-file=kaggle.json
    

Implementa Saxml

En esta sección, implementarás el servidor de administrador de Saxml, los servidores de modelos y el servidor HTTP.

Implementa el servidor de administrador de Saxml

  1. Crea el siguiente manifiesto saxml-admin-server.yaml:

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: sax-admin-server
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: sax-admin-server
      template:
        metadata:
          labels:
            app: sax-admin-server
        spec:
          hostNetwork: false
          containers:
          - name: sax-admin-server
            image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.2.0
            securityContext:
              privileged: true
            ports:
            - containerPort: 10000
            env:
            - name: GSBUCKET
              value: ADMIN_BUCKET_NAME

    Reemplaza ADMIN_BUCKET_NAME por el nombre del bucket que creaste en la sección Crea buckets de Cloud Storage. No incluyas el prefijo gs://.

  2. Aplica el manifiesto

    kubectl apply -f saxml-admin-server.yaml
    
  3. Verifica la implementación del servidor de administrador:

    kubectl get deployment
    

    El resultado es similar al siguiente:

    NAME                              READY   UP-TO-DATE   AVAILABLE   AGE
    sax-admin-server                  1/1     1            1           ##s
    

Implementa el servidor de modelos de Saxml

Sigue estas instrucciones para implementar el servidor de modelos para el modelo Gemma 2B o Gemma 7B.

Gemma 2B-it

  1. Crea el siguiente manifiesto saxml-model-server-1x1.yaml:

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: sax-model-server-v5e-1x1
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: gemma-server
      strategy:
        type: Recreate
      template:
        metadata:
          labels:
            app: gemma-server
            ai.gke.io/model: gemma-2b-it
            ai.gke.io/inference-server: saxml
            examples.ai.gke.io/source: user-guide
        spec:
          nodeSelector:
            cloud.google.com/gke-tpu-topology: 1x1
            cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
          hostNetwork: false
          restartPolicy: Always
          containers:
          - name: inference-server
            image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
            args:
            - "--jax_platforms=tpu"
            - "--platform_chip=tpuv5e"
            - "--platform_topology=1x1"
            - "--port=10001"
            - "--sax_cell=/sax/test"
            ports:
            - containerPort: 10001
            securityContext:
              privileged: true
            env:
            - name: SAX_ROOT
              value: "gs://ADMIN_BUCKET_NAME/sax-root"
            resources:
              requests:
                google.com/tpu: 1
              limits:
                google.com/tpu: 1

    Reemplaza ADMIN_BUCKET_NAME por el nombre del bucket que creaste en la sección Crea buckets de Cloud Storage. No incluyas el prefijo gs://.

  2. Aplica el manifiesto

    kubectl apply -f saxml-model-server-1x1.yaml
    
  3. Verifica el estado del Deployment del servidor de modelos:

    kubectl get deployment
    

    El resultado es similar al siguiente:

    NAME                                              READY   STATUS    RESTARTS   AGE
    sax-admin-server                                  1/1     Running   0          ##m
    sax-model-server-v5e-1x1                          1/1     Running   0          ##s
    

Gemma 7B-it

  1. Crea el siguiente manifiesto saxml-model-server-2x2.yaml:

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: sax-model-server-v5e-2x2
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: gemma-server
      strategy:
        type: Recreate
      template:
        metadata:
          labels:
            app: gemma-server
            ai.gke.io/model: gemma-7b-it
            ai.gke.io/inference-server: saxml
            examples.ai.gke.io/source: user-guide
        spec:
          nodeSelector:
            cloud.google.com/gke-tpu-topology: 2x2
            cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
          hostNetwork: false
          restartPolicy: Always
          containers:
          - name: inference-server
            image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
            args:
            - "--jax_platforms=tpu"
            - "--platform_chip=tpuv5e"
            - "--platform_topology=2x2"
            - "--port=10001"
            - "--sax_cell=/sax/test"
            ports:
            - containerPort: 10001
            securityContext:
              privileged: true
            env:
            - name: SAX_ROOT
              value: "gs://ADMIN_BUCKET_NAME/sax-root"
            resources:
              requests:
                google.com/tpu: 4
              limits:
                google.com/tpu: 4

    Reemplaza ADMIN_BUCKET_NAME por el nombre del bucket que creaste en la sección Crea buckets de Cloud Storage. No incluyas el prefijo gs://.

  2. Aplica el manifiesto

    kubectl apply -f saxml-model-server-2x2.yaml
    
  3. Verifica el estado del Deployment del servidor de modelos:

    kubectl get deployment
    

    El resultado es similar al siguiente:

    NAME                                              READY   STATUS    RESTARTS   AGE
    sax-admin-server                                  1/1     Running   0          ##m
    sax-model-server-v5e-2x2                          1/1     Running   0          ##s
    

Implementa el servidor HTTP de Saxml

En esta sección, implementarás el servidor HTTP de Saxml y crearás un servicio de IP de clúster que usarás para acceder al servidor.

  1. Crea el siguiente manifiesto saxml-http.yaml:

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: sax-http
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: sax-http
      template:
        metadata:
          labels:
            app: sax-http
        spec:
          hostNetwork: false
          containers:
          - name: sax-http
            image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.2.0
            imagePullPolicy: Always
            ports:
            - containerPort: 8888
            env:
            - name: SAX_ROOT
              value: "gs://ADMIN_BUCKET_NAME/sax-root"
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: sax-http-svc
    spec:
      selector:
        app: sax-http
      ports:
      - protocol: TCP
        port: 8888
        targetPort: 8888
      type: ClusterIP

    Reemplaza ADMIN_BUCKET_NAME por el nombre del bucket de Cloud Storage que almacena el servidor de administración de Saxml.

  2. Aplica el manifiesto

    kubectl apply -f saxml-http.yaml
    
  3. Verifica el estado de la implementación del servidor HTTP de Saxml:

    kubectl get deployment
    

    Gemma 2B-it

    El resultado es similar al siguiente:

      NAME                                              READY   STATUS    RESTARTS   AGE
      sax-admin-server                                  1/1     Running   0          ##m
      sax-model-server-v5e-1x1                          1/1     Running   0          ##m
      sax-http                                          1/1     Running   0          ##s
    

    Gemma 7B-it

    El resultado es similar al siguiente:

      NAME                                              READY   STATUS    RESTARTS   AGE
      sax-admin-server                                  1/1     Running   0          ##m
      sax-model-server-v5e-2x2                          1/1     Running   0          ##m
      sax-http                                          1/1     Running   0          ##s
    

Descarga el punto de control del modelo

En esta sección, ejecutas un trabajo de Kubernetes que recupera, descarga y almacena el punto de control del modelo. Sigue los pasos para el modelo de Gemma que deseas usar:

Gemma 2B-it

  1. Crea el siguiente manifiesto job-2b.yaml:

    apiVersion: v1
    kind: ConfigMap
    metadata:
      name: fetch-model-scripts
    data:
      fetch_model.sh: |-
        #!/usr/bin/bash -x
        pip install kaggle --break-system-packages && \
    
        MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
        VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \
    
        mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
        kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
        echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \
    
        gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
        echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: data-loader-2b
      labels:
        app: data-loader-2b
    spec:
      ttlSecondsAfterFinished: 120
      template:
        metadata:
          labels:
            app: data-loader-2b
        spec:
          restartPolicy: OnFailure
          containers:
          - name: gcloud
            image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
            command:
            - /scripts/fetch_model.sh
            env:
            - name: BUCKET_NAME
              value: CHECKPOINTS_BUCKET_NAME
            - name: KAGGLE_CONFIG_DIR
              value: /kaggle
            - name: MODEL_PATH
              value: "google/gemma/pax/2b-it/2"
            volumeMounts:
            - mountPath: "/kaggle/"
              name: kaggle-credentials
              readOnly: true
            - mountPath: "/scripts/"
              name: scripts-volume
              readOnly: true
          volumes:
          - name: kaggle-credentials
            secret:
              defaultMode: 0400
              secretName: kaggle-secret
          - name: scripts-volume
            configMap:
              defaultMode: 0700
              name: fetch-model-scripts

    Reemplaza CHECKPOINTS_BUCKET_NAME por el nombre del bucket que creaste en la sección Crea buckets de Cloud Storage. No incluyas el prefijo gs://.

  2. Aplica el manifiesto

    kubectl apply -f job-2b.yaml
    
  3. Espera a que se complete el trabajo:

    kubectl wait --for=condition=complete --timeout=180s job/data-loader-2b
    

    El resultado es similar al siguiente:

    job.batch/data-loader-2b condition met
    
  4. Verifica que el trabajo se haya completado de forma correcta:

    kubectl get job/data-loader-2b
    

    El resultado es similar al siguiente:

    NAME             COMPLETIONS   DURATION   AGE
    data-loader-2b   1/1           ##s        #m##s
    
  5. Visualiza los registros del trabajo:

    kubectl logs --follow job/data-loader-2b
    

El trabajo sube el punto de control a gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000.

Gemma 7B-it

  1. Crea el siguiente manifiesto job-7b.yaml:

    apiVersion: v1
    kind: ConfigMap
    metadata:
      name: fetch-model-scripts
    data:
      fetch_model.sh: |-
        #!/usr/bin/bash -x
        pip install kaggle --break-system-packages && \
    
        MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
        VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \
    
        mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
        kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
        echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \
    
        gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
        echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: data-loader-7b
      labels:
        app: data-loader-7b
    spec:
      ttlSecondsAfterFinished: 120
      template:
        metadata:
          labels:
            app: data-loader-7b
        spec:
          restartPolicy: OnFailure
          containers:
          - name: gcloud
            image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
            command:
            - /scripts/fetch_model.sh
            env:
            - name: BUCKET_NAME
              value: CHECKPOINTS_BUCKET_NAME
            - name: KAGGLE_CONFIG_DIR
              value: /kaggle
            - name: MODEL_PATH
              value: "google/gemma/pax/7b-it/2"
            volumeMounts:
            - mountPath: "/kaggle/"
              name: kaggle-credentials
              readOnly: true
            - mountPath: "/scripts/"
              name: scripts-volume
              readOnly: true
          volumes:
          - name: kaggle-credentials
            secret:
              defaultMode: 0400
              secretName: kaggle-secret
          - name: scripts-volume
            configMap:
              defaultMode: 0700
              name: fetch-model-scripts

    Reemplaza CHECKPOINTS_BUCKET_NAME por el nombre del bucket que creaste en la sección Crea buckets de Cloud Storage. No incluyas el prefijo gs://.

  2. Aplica el manifiesto

    kubectl apply -f job-7b.yaml
    
  3. Espera a que se complete el trabajo:

    kubectl wait --for=condition=complete --timeout=360s job/data-loader-7b
    

    El resultado es similar al siguiente:

    job.batch/data-loader-7b condition met
    
  4. Verifica que el trabajo se haya completado de forma correcta:

    kubectl get job/data-loader-7b
    

    El resultado es similar al siguiente:

    NAME             COMPLETIONS   DURATION   AGE
    data-loader-7b   1/1           ##s        #m##s
    
  5. Visualiza los registros del trabajo:

    kubectl logs --follow job/data-loader-7b
    

El trabajo sube el punto de control a gs://CHECKPOINTS_BUCKET_NAME/gemma_7b_it/checkpoint_00000000.

Expón el servidor HTTP de Saxml

Puedes acceder al servidor HTTP de Saxml a través del Service de ClusterIP que creaste cuando implementaste el servidor HTTP de Saxml. Solo se puede acceder a los servicios de ClusterIP desde el clúster. Por lo tanto, para acceder al servicio desde fuera del clúster, completa los siguientes pasos:

  1. Establece una sesión de redirección de puertos:

    kubectl port-forward service/sax-http-svc 8888:8888
    
  2. Para comprobar que puedes acceder al servidor HTTP de Saxml, abre una terminal nueva y ejecuta el siguiente comando:

    curl -s localhost:8888
    

    El resultado es similar al siguiente:

    {
        "Message": "HTTP Server for SAX Client"
    }
    

El servidor HTTP de Saxml encapsula la interfaz del cliente al sistema Saxml y la expone a través de un conjunto de APIs de REST. Usas estas APIs para publicar, administrar y también interactuar con los modelos Gemma 2B y Gemma 7B.

Publica el modelo de Gemma

A continuación, puedes publicar el modelo de Gemma en un servidor de modelos que se ejecute en un grupo de nodos de porción de TPU. Usa la API publish del servidor HTTP de Saxml para publicar un modelo. Sigue estos pasos para publicar el modelo de parámetros Gemma 2B o 7B.

Para obtener más información sobre la API del servidor HTTP de Saxml, consulta APIs de HTTP de Saxml.

Gemma 2B-it

  1. Asegúrate de que tu sesión de redirección de puertos aún esté activa:

    curl -s localhost:8888
    
  2. Publica el parámetro de Gemma 2B:

    curl --request POST \
    --header "Content-type: application/json" \
    -s \
    localhost:8888/publish \
    --data \
    '{
        "model": "/sax/test/gemma2bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
        "replicas": "1"
    }'
    

    El resultado es similar al siguiente:

    {
        "model": "/sax/test/gemma2bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
        "replicas": 1
    }
    

    Consulta el siguiente paso para supervisar el progreso de la implementación.

  3. Para supervisar el progreso, observa los registros en un Pod de servidor de modelos de la implementación de sax-model-server-v5e-1x1.

    kubectl logs --follow deployment/sax-model-server-v5e-1x1
    

    Esta implementación puede tardar hasta cinco minutos en completarse. Espera hasta que veas un mensaje similar al siguiente:

    I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
    I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma2bfp16
    
  4. Verifica que puedas acceder al modelo si muestras la información del modelo:

    curl --request GET \
    --header "Content-type: application/json" \
    -s \
    localhost:8888/listcell \
    --data \
    '{
        "model": "/sax/test/gemma2bfp16"
    }'
    

    El resultado es similar al siguiente:

    {
        "model": "/sax/test/gemma2bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
        "max_replicas": 1,
        "active_replicas": 1
    }
    

Gemma 7B-it

  1. Asegúrate de que tu sesión de redirección de puertos aún esté activa:

    curl -s localhost:8888
    
  2. Publica el parámetro de Gemma 7B:

    curl --request POST \
    --header "Content-type: application/json" \
    -s \
    localhost:8888/publish \
    --data \
    '{
        "model": "/sax/test/gemma7bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
        "replicas": "1"
    }'
    

    El resultado es similar al siguiente:

    {
        "model": "/sax/test/gemma7bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
        "replicas": 1
    }
    

    Consulta el siguiente paso para supervisar el progreso de la implementación.

  3. Para supervisar el progreso, observa los registros en un Pod de servidor de modelos de la implementación de sax-model-server-v5e-2x2.

    kubectl logs --follow deployment/sax-model-server-v5e-2x2
    

    Espera hasta que veas un mensaje similar al siguiente:

    I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
    I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma7bfp16
    
  4. Muestra la información del modelo para verificar que el modelo se publicó:

    curl --request GET \
    --header "Content-type: application/json" \
    -s \
    localhost:8888/listcell \
    --data \
    '{
        "model": "/sax/test/gemma7bfp16"
    }'
    

    El resultado es similar al siguiente:

    {
        "model": "/sax/test/gemma7bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
        "max_replicas": 1,
        "active_replicas": 1
    }
    

Usa el modelo

Puedes interactuar con los modelos Gemma 2B o 7B. Usa la API generate del servidor HTTP de Saxml para enviar un mensaje al modelo.

Gemma 2B-it

Entrega una solicitud de instrucción con el extremo generate del servidor HTTP de Saxml:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/generate \
--data \
'{
  "model": "/sax/test/gemma2bfp16",
  "query": "What are the top 5 most popular programming languages?"
}'

El siguiente es un ejemplo de la respuesta del modelo. El resultado real varía según el mensaje que entregues:

[
    [
        "\n\n1. **Python**\n2. **JavaScript**\n3. **Java**\n4. **C++**\n5. **Go**",
        -3.0704939365386963
    ]
]

Puedes ejecutar el comando con diferentes parámetros de query. También puedes modificar parámetros adicionales, como temperature, top_k, topc_p, mediante la API de generate. Para obtener más información sobre la API del servidor HTTP de Saxml, consulta APIs de HTTP de Saxml.

Gemma 7B-it

Entrega una solicitud de instrucción con el extremo generate del servidor HTTP de Saxml:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/generate \
--data \
'{
  "model": "/sax/test/gemma7bfp16",
  "query": "What are the top 5 most popular programming languages?"
}'

El siguiente es un ejemplo de la respuesta del modelo. El resultado puede variar en cada mensaje que entregues:

[
    [
        "\n\n**1. JavaScript**\n\n* Most widely used language on the web.\n* Used for front-end development, such as websites and mobile apps.\n* Extensive libraries and frameworks available.\n\n**2. Python**\n\n* Known for its simplicity and readability.\n* Versatile, used for various tasks, including data science, machine learning, and web development.\n* Large and active community.\n\n**3. Java**\n\n* Object-oriented language widely used in enterprise applications.\n* Used for web applications, mobile apps, and enterprise software.\n* Strong ecosystem and support.\n\n**4. Go**\n\n",
        -16.806324005126953
    ]
]

Puedes ejecutar el comando con diferentes parámetros de query. También puedes modificar parámetros adicionales, como temperature, top_k, topc_p, mediante la API de generate. Para obtener más información sobre la API del servidor HTTP de Saxml, consulta APIs de HTTP de Saxml.

Anula la publicación del modelo

Sigue estos pasos para anular la publicación de tu modelo:

Gemma 2B-it

Para anular la publicación del modelo de Gemma 2B-it, ejecuta el siguiente comando:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/unpublish \
--data \
'{
    "model": "/sax/test/gemma2bfp16"
}'

El resultado es similar al siguiente:

{
    "model": "/sax/test/gemma2bfp16"
}

Puedes ejecutar el comando con diferentes mensajes que se pasan en el parámetro query.

Gemma 7B-it

Para anular la publicación del modelo de Gemma 7B-it, ejecuta el siguiente comando:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/unpublish \
--data \
'{
    "model": "/sax/test/gemma7bfp16"
}'

El resultado es similar al siguiente:

{
    "model": "/sax/test/gemma7bfp16"
}

Puedes ejecutar el comando con diferentes mensajes que se pasan en el parámetro query.

Soluciona problemas

Realiza una limpieza

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=${LOCATION}
gcloud iam service-accounts delete --quiet wi-sax@${PROJECT_ID}.iam.gserviceaccount.com
gcloud storage rm --recursive gs://ADMIN_BUCKET_NAME
gcloud storage rm --recursive gs://CHECKPOINTS_BUCKET_NAME

Reemplaza lo siguiente:

  • ADMIN_BUCKET_NAME: el nombre del bucket de Cloud Storage que almacena el servidor de administrador de Saxml.
  • CHECKPOINTS_BUCKET_NAME: es el nombre del bucket de Cloud Storage que almacena los puntos de control del modelo.

¿Qué sigue?