Entrega modelos abiertos de Gemma con TPU en GKE con Saxml

Estándar

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:

Obtén información sobre la disponibilidad actual de la versión de TPU con la arquitectura del sistema de Cloud TPU.
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:

Prepara un clúster de GKE Standard con la topología de TPU recomendada según las características del modelo.
Implementa componentes de Saxml en GKE.
Obtén y publica el modelo de parámetros Gemma 2B o Gemma 7B.
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.

En la consola de Google Cloud, inicia una instancia de Cloud Shell:
Abrir Cloud Shell
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 ID del proyecto de Google Cloud.
- LOCATION: es el nombre de la zona de Compute Engine en la que los tipos de máquina de TPU v5e están disponibles.

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:

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.

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:

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.

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:

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

Configura kubectl para comunicarse con tu clúster:

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}

Crea una cuenta de servicio de Kubernetes para que tu aplicación use:
```
gcloud iam service-accounts create wi-sax
```

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

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]"
```

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:

Accede a la página de consentimiento del modelo en Kaggle.com.
Ingresa en Kaggle si aún no lo has hecho.
Haz clic en Solicitar acceso.
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.
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:

En el navegador, ve a la configuración de Kaggle.
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:

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

Crea un secreto de Kubernetes para las credenciales de Kaggle

En Cloud Shell, haz lo siguiente:

Configura kubectl para comunicarse con tu clúster:

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}

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

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://.

Aplica el manifiesto

kubectl apply -f saxml-admin-server.yaml

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

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://.

Aplica el manifiesto

kubectl apply -f saxml-model-server-1x1.yaml

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

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://.

Aplica el manifiesto

kubectl apply -f saxml-model-server-2x2.yaml

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.

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.

Aplica el manifiesto
```
kubectl apply -f saxml-http.yaml
```

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

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://.

Aplica el manifiesto
```
kubectl apply -f job-2b.yaml
```

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

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

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

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://.

Aplica el manifiesto
```
kubectl apply -f job-7b.yaml
```

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

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

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:

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

kubectl port-forward service/sax-http-svc 8888:8888

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

Asegúrate de que tu sesión de redirección de puertos aún esté activa:
```
curl -s localhost:8888
```

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.

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

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

Asegúrate de que tu sesión de redirección de puertos aún esté activa:
```
curl -s localhost:8888
```

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.

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

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
    ]
]

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

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 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?

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.

Entrega modelos abiertos de Gemma con TPU en GKE con Saxml

Antecedentes

Gemma

TPU

Saxml

Objetivos

Arquitectura

Antes de comenzar

Check for the roles

Grant the roles

Prepara el entorno para Gemma

Inicia Cloud Shell

Crea un clúster de GKE Standard

Gemma 2B-it

Gemma 7B-it

Crea los buckets de Cloud Storage

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

Obtén acceso al modelo

Firma el acuerdo de consentimiento de licencia

Genera un token de acceso

Sube el token de acceso a Cloud Shell

Crea un secreto de Kubernetes para las credenciales de Kaggle

Implementa Saxml

Implementa el servidor de administrador de Saxml

Implementa el servidor de modelos de Saxml

Gemma 2B-it

Gemma 7B-it

Implementa el servidor HTTP de Saxml

Gemma 2B-it

Gemma 7B-it

Descarga el punto de control del modelo

Gemma 2B-it

Gemma 7B-it

Expón el servidor HTTP de Saxml

Publica el modelo de Gemma

Gemma 2B-it

Gemma 7B-it

Usa el modelo

Gemma 2B-it

Gemma 7B-it

Anula la publicación del modelo

Gemma 2B-it

Gemma 7B-it

Soluciona problemas

Realiza una limpieza

Borra los recursos implementados

¿Qué sigue?