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 características 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.

Formación

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 livianos y disponibles a nivel abierto que se lanzaron 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 para 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 marcos de trabajo 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 más información sobre las TPU en GKE.

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

Gemma 2B: Modelo ajustado de instrucciones alojado en un grupo de nodos TPU v5e con topología 1x1 que representa un chip TPU. El tipo de máquina para los nodos es ct5lp-hightpu-1t.
Gemma 7B: Modelo ajustado a las instrucciones alojado en un grupo de nodos de TPU v5e con 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 entrega modelos Paxml, JAX y PyTorch para la inferencia. El sistema Saxml consta de los siguientes componentes:

Celda de Saxml o clúster Sax: compuesto por 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 que entreguen y ayuda a los clientes a ubicar servidores de modelos que entregan 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 a las APIs de REST para interactuar con el sistema Saxml, incluidos los extremos para publicar, enumerar y anular la publicación de modelos, y generar predicciones.

Objetivos

Este instructivo está dirigido a clientes de IA generativa que usan JAX, usuarios nuevos o existentes de GKE, ingenieros de AA, ingenieros de MLOps (DevOps) o administradores de plataformas interesados en usar las funciones de organización de contenedores de Kubernetes para entregando a Gemma.

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: configurado con una TPU v5e con una topología 1x1. Una instancia del servidor de modelos de Saxml está configurada para usar este grupo de nodos.
- Gemma 7B: Configurado con una TPU v5e con una topología 2x2. Una instancia del servidor de modelos de Saxml está configurada para usar este grupo de nodos.
Un grupo de nodos de CPU predeterminado en el que se implementan el servidor administrador Saxml y el servidor HTTP 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 Gemma.

Esta arquitectura tiene las siguientes características:

Un Artifact Registry público administra las imágenes de contenedores 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 que se crean 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.

En la página del selector de proyectos de la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.

Ir al selector de proyectos

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

Habilita la API necesaria.

Habilita la API

En la página del selector de proyectos de la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.

Ir al selector de proyectos

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

Habilita la API necesaria.

Habilita la API

Asegúrate de tener los siguientes roles en el proyecto: roles/container.admin, roles/iam.serviceAccountAdmin
Verifica los roles
1. En la consola de Google Cloud, ve a la página IAM.
  Ir a IAM
2. Selecciona el proyecto.
3. En la columna Principal, busca la fila que tiene tu dirección de correo electrónico.
  
  Si tu dirección de correo electrónico no está en esa columna, no tienes ningún rol.
4. En la columna Función de la fila con la dirección de correo electrónico, verifica si la lista de roles incluye los roles necesarios.
Otorga los roles
1. En la consola de Google Cloud, ve a la página IAM.
  Ir a IAM
2. Selecciona el proyecto.
3. Haz clic en Grant access.
4. En el campo Principales nuevas, ingresa tu dirección de correo electrónico.
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 a pedido.
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: El nombre de la zona de Compute Engine en la que los tipos de máquinas 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 Standard 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}

Entregas el modelo Gemma 2B en este grupo de nodos.

Gemma 7B-it

Usa Cloud Shell para realizar las siguientes acciones:

Crea un clúster Standard 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}

Entregas el modelo Gemma 7B en este grupo de nodos.

Crea los buckets de Cloud Storage

Crear dos bucket de Cloud Storage para administrar el estado del servidor Saxml Admin y los puntos de control del modelo

En Cloud Shell, ejecuta lo siguiente:

Crea un bucket de Cloud Storage para almacenar configuraciones del servidor 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.
Crea un bucket de Cloud Storage para almacenar 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.

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 secreto 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.
Accede a 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.

Se descargará un archivo con el nombre 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.
Selecciona 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 administración de Saxml, los servidores de modelos y el servidor HTTP.

Implementa el servidor de administración 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 debería ser similar al siguiente ejemplo:

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

Implementa el servidor de modelos Saxml

Sigue estas instrucciones a fin de implementar el servidor de modelos para los modelos 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 debería ser similar al siguiente ejemplo:

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 debería ser similar al siguiente ejemplo:

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 administrador 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 debería ser similar al siguiente ejemplo:

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 debería ser similar al siguiente ejemplo:

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. Sigue los pasos según 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

Los puntos de control se suben 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. Incluye 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

Los puntos de control se suben 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 de un Service de ClusterIP que creaste en el paso anterior. 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

Verifica que puedas acceder al servidor HTTP de Saxml a través de la apertura de una terminal nueva y la ejecución del 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. Usarás estas APIs para publicar, administrar y también interactuar con los modelos Gemma 2B y Gemma 7B.

Publica el modelo Gemma

Después de configurar los componentes de Saxml, puedes publicar el modelo Gemma en un servidor de modelos que se ejecute en un grupo de nodos de porción de TPU. Usa la API de publish del servidor HTTP de Saxml para publicar un modelo. Para obtener más información sobre la API del servidor HTTP de Saxml, consulta las APIs de HTTP de Saxml. Sigue estos pasos para publicar el modelo de parámetros Gemma 2B o 7B.

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
}

La implementación del modelo puede tardar unos minutos.

Para supervisar el progreso, observa los registros en un Pod de servidor de modelos de la implementación 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 ver 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

Para verificar que puedas acceder al modelo, muestra 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 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
}

La implementación del modelo puede tardar unos minutos.

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

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

Espera hasta ver 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 se haya publicado:

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 prompts con el extremo generate del servidor HTTP 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 puede variar en cada mensaje que entregues:

[
    [
        "\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, 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 prompts con el extremo generate del servidor HTTP 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 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 será similar al siguiente ejemplo:

{
    "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 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 será similar al siguiente ejemplo:

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

Limpia

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 administración de Saxml.
CHECKPOINTS_BUCKET_NAME: 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.
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

Formación

Gemma

TPU

Saxml

Objetivos

Arquitectura

Antes de comenzar

Verifica los roles

Otorga los 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 contrato 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 administración de Saxml

Implementa el servidor de modelos 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 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

Limpia

Borra los recursos implementados

¿Qué sigue?