Entrega modelos abiertos de Gemma con TPU en GKE con Saxml

Estándar

---

Precaución: A partir del 24 de abril de 2025, se reducirá la prioridad de Saxml en GKE. Esto significa que el proyecto no recibirá más actualizaciones. Las implementaciones existentes de Saxml seguirán funcionando como de costumbre sin interrupciones. Te recomendamos enfáticamente que migres a JetStream, el framework de inferencia de código abierto actualizado de Google para la entrega de LLM de alto rendimiento en TPU y GPU. JetStream ofrece procesamiento por lotes y cuantización continuos para mejorar la capacidad de procesamiento y la eficiencia de la memoria. Para ver un ejemplo de migración, consulta Entrega Gemma con TPU en GKE con JetStream.

En este instructivo, se muestra cómo implementar y entregar un modelo de lenguaje grande (LLM) de Gemma 2 con TPU en GKE con el framework de entrega Saxml. En este instructivo, se proporcionan los conceptos básicos para comprender y explorar la implementación práctica de LLM para la inferencia en un entorno de Kubernetes administrado. Implementas un contenedor prediseñado con Saxml en GKE. También configurarás GKE para que cargue las ponderaciones de Gemma 2B y 7B desde Cloud Storage en el tiempo de ejecución.

Este instructivo está dirigido a ingenieros de aprendizaje automático (AA), administradores y operadores de plataformas, y especialistas en datos y en IA que estén interesados en usar las capacidades de organización de contenedores de Kubernetes para entregar LLMs. Para obtener más información sobre los roles comunes y las tareas de ejemplo a las que hacemos referencia en el contenido deGoogle Cloud , consulta Roles de usuario y tareas comunes de GKE Enterprise.

Antes de leer esta página, asegúrate de estar familiarizado con los siguientes temas:

• Disponibilidad de la versión actual de TPU con la arquitectura del sistema de Cloud TPU
• TPUs en GKE

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.

Fondo

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, pero también puedes ajustar estos modelos en el caso de tareas especializadas.

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

TPU

Las TPU 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.

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: Es un modelo ajustado según instrucciones alojado en un grupo de nodos de 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: Es un modelo ajustado según instrucciones alojado en un grupo de nodos de 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 Saxml o clúster de Sax: Es un servidor de administrador 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 que los publiquen y ayuda a los clientes a ubicar los servidores de modelos que publican modelos 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 APIs de REST para interactuar con el sistema Saxml. Las APIs de REST incluyen extremos para publicar, enumerar y anular la publicación de modelos, y generar predicciones.

Objetivos

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:

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 configura una instancia del servidor de modelos de Saxml para usar este grupo de nodos.
  • Gemma 7B: Se configuró con una TPU v5e con una topología 2x2. Se configura una instancia del servidor de modelos de Saxml para usar este grupo de nodos.
• Es 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 administrador.
  • 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 de 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 los que se crean en este instructivo.

Antes de comenzar

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

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

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the required API.

Enable the API

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

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the required API.

Enable the API

• Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin, roles/iam.policyAdmin, roles/iam.securityAdmin, roles/iam.roleAdmin

  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 Otorgar acceso.

  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 uno.
  6. Para otorgar roles adicionales, haz clic en Agregar otro rol y agrega uno más.
  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 Google Cloud consola, 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:

     • PROJECT_ID: El Google Cloud ID del proyecto.
     • LOCATION: Es el nombre de la zona de Compute Engine en la que están disponibles los tipos de máquinas de TPU v5e.

  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}
     

     En este grupo de nodos, se entrega el modelo Gemma 2B.

  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}
     

     En este grupo de nodos, se entrega el modelo Gemma 7B.

  Crea los buckets de Cloud Storage

  Crea dos bucket 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 las opciones de configuración del servidor de administrador 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 administrador de Saxml.

  2. Crea un bucket de Cloud Storage para almacenar los puntos de control del modelo:

     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.

  Firma el acuerdo de consentimiento de licencia

  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 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 administración de Saxml, los servidores de modelos y el servidor HTTP. En este instructivo, se usan manifiestos de implementación de Kubernetes. Un Deployment es un objeto de la API de Kubernetes que te permite ejecutar varias réplicas de Pods que se distribuyen entre los nodos de un clúster.

  Implementa el servidor de administración de Saxml

  En esta sección, implementarás 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 administración:

     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 del 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 de la 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
     

  Completado correctamente: Implementaste un servidor de modelos en un grupo de nodos de TPU v5e con una porción de 1 x 1 para entregar el modelo Gemma 2B.

  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 de la 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
     

  Completado correctamente: Implementaste un servidor de modelos en un grupo de nodos de TPU v5e con una porción 2x2 para entregar el modelo Gemma 7B.

  Implementa el servidor HTTP de Saxml

  En esta sección, implementarás el servidor HTTP de Saxml y crearás un servicio de IP del 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 administrador 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, ejecutarás un trabajo de Kubernetes que recupera, descarga y almacena el punto de control del modelo. Un controlador de Job en Kubernetes crea uno o más Pods y garantiza que ejecuten correctamente una tarea específica.

  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 correctamente:

     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 objeto Job 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. Incluye 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 correctamente:

     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 objeto Job 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 servicio 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. Usas la API de publish del servidor HTTP de Saxml para publicar un modelo. Sigue estos pasos para publicar el modelo de parámetros de 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 siga activa:

     curl -s localhost:8888
     

  2. Publica el parámetro 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. Supervisa el progreso observando los registros en un Pod del 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 puedes acceder al modelo mostrando su información:

     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
     }
     

  Completado correctamente: Publicaste un modelo de Gemma 2B. Ahora puedes interactuar con el modelo.

  Gemma 7B-it

  1. Asegúrate de que tu sesión de redirección de puertos siga activa:

     curl -s localhost:8888
     

  2. Publica el parámetro 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. Supervisa el progreso observando los registros en un Pod del 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. Verifica que el modelo se haya publicado mostrando su información:

     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
     }
     

  Completado correctamente: Publicaste un modelo de Gemma 7B. Ahora puedes interactuar con el modelo.

  Usa el modelo

  Puedes interactuar con los modelos Gemma 2B o 7B. Usa la API de generate del servidor HTTP de Saxml para enviar una instrucción 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?"
  }'
  

  A continuación, se muestra un ejemplo de la respuesta del modelo. El resultado real varía según la instrucción que proporciones:

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

  Puedes ejecutar el comando con diferentes parámetros query. También puedes modificar parámetros adicionales, como temperature, top_k y topc_p, con 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?"
  }'
  

  A continuación, se muestra un ejemplo de la respuesta del modelo. El resultado puede variar en cada instrucción que proporciones:

  [
      [
          "\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 query. También puedes modificar parámetros adicionales, como temperature, top_k y topc_p, con 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 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 instrucciones que se pasan en el parámetro query.

  Gemma 7B-it

  Para anular la publicación del modelo 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 instrucciones que se pasan en el parámetro query.

  Soluciona problemas

  • Si recibes el mensaje Empty reply from server, es posible que el contenedor no haya terminado de descargar los datos del modelo. Vuelve a verificar los registros del Pod en busca del mensaje Connected, que indica que el modelo está listo para entregar.
  • Si ves Connection refused, verifica que tu redirección de puertos esté activa.

  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 apliquen cargos a 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: Es el nombre del bucket de Cloud Storage que almacena el servidor de administración de Saxml.
  • CHECKPOINTS_BUCKET_NAME: Es el nombre del bucket de Cloud Storage que almacena los puntos de control del modelo.

  ¿Qué sigue?

  • Obtén más información sobre las TPU en GKE.
  • Explora el repositorio de GitHub de Saxml, incluidas las APIs de HTTP de Saxml.
  • Explora Vertex AI Model Garden.
  • Descubre cómo ejecutar cargas de trabajo de IA/AA optimizadas con las capacidades de organización de la plataforma de GKE.
  • Explora arquitecturas de referencia, diagramas y prácticas recomendadas sobre Google Cloud. Consulta nuestro Cloud Architecture Center.