Entrega modelos abiertos de Gemma con TPU en Vertex AI Prediction con Saxml

En esta guía, se muestra cómo entregar un modelo de lenguaje grande (LLM) de modelos abiertos de Gemma con unidades de procesamiento tensorial (TPU) en Vertex AI Prediction con Saxml. En esta guía, descargarás los modelos de Gemma ajustados a las instrucciones de los parámetros 2B y 7B en Cloud Storage y los implementarás en Vertex AI Prediction que ejecuta Saxml en TPU.

Formación

Mediante la entrega de Gemma con TPU en Vertex AI Prediction con Saxml. Puedes aprovechar una solución de IA administrada que se encarga de la infraestructura de bajo nivel y ofrece una manera rentable de entregar LLM. En esta sección, se describen las tecnologías clave que se usan en este instructivo.

Gemma

Gemma es un conjunto de modelos de inteligencia artificial (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 para tareas especializadas.

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

Saxml

Saxml es un sistema experimental que entrega modelos Paxml, JAX y PyTorch para la inferencia. A los fines de este instructivo, abordaremos cómo entregar Gemma en TPU que son más rentables para Saxml. La configuración de las GPU es similar. Saxml ofrece secuencias de comandos a fin de compilar contenedores para Vertex AI Prediction que usaremos en este instructivo.

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.

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

Gemma 2B: Alojado en un grupo de nodos de 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: 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.

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 las API de Vertex AI API and Artifact Registry API.

Habilita las 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 las API de Vertex AI API and Artifact Registry API.

Habilita las API

En la consola de Google Cloud, activa Cloud Shell.

Activar Cloud Shell

En la parte inferior de la consola de Google Cloud, se inicia una sesión de Cloud Shell en la que se muestra una ventana de línea de comandos. Cloud Shell es un entorno de shell con Google Cloud CLI ya instalada y con valores ya establecidos para el proyecto actual. La sesión puede tardar unos segundos en inicializarse.

En este instructivo, se supone que usas Cloud Shell para interactuar con Google Cloud. Si quieres usar un shell diferente en lugar de Cloud Shell, realiza la siguiente configuración adicional:

Instala Google Cloud CLI.
Para inicializar la CLI de gcloud, ejecuta el siguiente comando:
```
gcloud init
```

Sigue la documentación de Artifact Registry para instalar Docker.
Asegúrate de tener cuotas suficientes para 5 chips TPU v5e para Vertex AI Prediction.
Crea una cuenta de Kaggle, si todavía no tienes una.

Obtén acceso al modelo

Para obtener acceso a los modelos de Gemma para la implementación en Vertex AI Prediction, 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.
Si aún no lo has hecho, accede a Kaggle.
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 el bucket de Cloud Storage

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

En Cloud Shell, ejecuta lo siguiente:

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.

Copia el modelo al bucket de Cloud Storage

En Cloud Shell, ejecuta lo siguiente:

pip install kaggle --break-system-packages

# For Gemma 2B
mkdir -p /data/gemma_2b-it
kaggle models instances versions download google/gemma/pax/2b-it/1 --untar -p /data/gemma_2b-it
gsutil -m cp -R /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/

# For Gemma 7B
mkdir -p /data/gemma_7b-it
kaggle models instances versions download google/gemma/pax/7b-it/1 --untar -p /data/gemma_7b-it
gsutil -m cp -R /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/

Crea un repositorio de Artifact Registry

Crea un repositorio de Artifact Registry para almacenar la imagen de contenedor que crearás en la próxima sección.

Habilita el servicio de la API de Artifact Registry para tu proyecto.

gcloud services enable artifactregistry.googleapis.com

Ejecuta el siguiente comando en la shell para crear un repositorio de Artifact Registry:

gcloud artifacts repositories create saxml \
 --repository-format=docker \
 --location=LOCATION \
 --description="Saxml Docker repository"

Reemplaza LOCATION por la región en la que Artifact Registry almacena la imagen de contenedor. Más adelante, debes crear un recurso de modelo de Vertex AI en un extremo regional que coincida con esta región, por lo que debes elegir una región en la que Vertex AI tenga un extremo regional, como us-west1 para las TPU.

Envía la imagen del contenedor a Artifact Registry

El contenedor Saxml ya compilado está disponible en us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest. Cópialo en tu Artifact Registry. Configurar Docker para acceder a Artifact Registry Luego, envía tu imagen de contenedor a tu repositorio de Artifact Registry.

Si deseas permitir que tu instalación local de Docker se envíe a Artifact Registry en la región que seleccionaste, ejecuta el siguiente comando en tu shell:
```
gcloud auth configure-docker LOCATION-docker.pkg.dev
```
- Reemplaza LOCATION por la región en la que creaste tu repositorio.

Para copiar la imagen de contenedor que acabas de compilar a Artifact Registry, ejecuta el siguiente comando en tu shell:

docker tag us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest

Para enviar la imagen de contenedor que acabas de compilar a Artifact Registry, ejecuta el siguiente comando en tu shell:
```
docker push LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest
```
Reemplaza lo siguiente, tal como lo hiciste en la sección anterior:
- LOCATION: Es la región de tu repositorio de Artifact Registry.
- PROJECT_ID: el ID del proyecto de Google Cloud

Implementa el modelo

Sube un modelo

Para subir un recurso Model que use tu contenedor de Saxml, ejecuta el siguiente comando gcloud ai models upload:

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma2BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Gemma 7B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma7BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Reemplaza lo siguiente:

PROJECT_ID: el ID del proyecto de Google Cloud
LOCATION_ID: la región en la que usas Vertex AI. Ten en cuenta que las TPU solo están disponibles en us-west1.
DEPLOYED_MODEL_NAME: Un nombre para DeployedModel También puedes usar el nombre comercial de Model para DeployedModel.

Crea un extremo

Debes implementar el modelo en un extremo antes de que pueda usarse para entregar predicciones en línea. Si implementas un modelo en un extremo existente, puedes omitir este paso. En el siguiente ejemplo, se usa el comando gcloud ai endpoints create:

gcloud ai endpoints create \
  --region=LOCATION \
  --display-name=ENDPOINT_NAME

Reemplaza lo siguiente:

LOCATION_ID: la región en la que usas Vertex AI.
ENDPOINT_NAME: el nombre visible para el extremo.

La herramienta de la CLI de Google Cloud puede tardar unos segundos en crear el extremo.

Implementa el modelo en el extremo

Una vez que el extremo esté listo, implementa el modelo en el extremo.

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai endpoints deploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --model=$MODEL_ID \
  --display-name=DEPLOYED_MODEL_NAME \
  --machine-type=ct5lp-hightpu-4t \
  --traffic-split=0=100

Reemplaza lo siguiente:

LOCATION_ID: la región en la que usas Vertex AI.
ENDPOINT_NAME: el nombre visible para el extremo.
DEPLOYED_MODEL_NAME: Un nombre para DeployedModel También puedes usar el nombre comercial de Model para DeployedModel.

Gemma 2B se puede implementar en una máquina ct5lp-hightpu-1t más pequeña; en ese caso, debes especificar --platform_topology=1x1 cuando subas el modelo.

La herramienta de Google Cloud CLI puede tardar unos minutos en implementar el modelo en el extremo. Cuando el modelo se implementa de forma correcta, este comando imprime el siguiente resultado:

  Deployed a model to the endpoint xxxxx. Id of the deployed model: xxxxx.

Obtén predicciones en línea del modelo implementado

Para invocar el modelo a través del extremo de Vertex AI Prediction, formatea la solicitud de predicción con un objeto JSON de solicitud de inferencia estándar .

En el siguiente ejemplo, se usa el comando gcloud ai endpoints predict:

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

gcloud ai endpoints predict $ENDPOINT_ID \
  --region=LOCATION \
  --http-headers=Content-Type=application/json \
  --json-request instances.json

Reemplaza lo siguiente:

LOCATION_ID: la región en la que usas Vertex AI.
ENDPOINT_NAME: el nombre visible para el extremo.
instances.json tiene el siguiente formato: {"instances": [{"text_batch": "<your prompt>"},{...}]}

Realice una limpieza

Para evitar incurrir en más cargos de Vertex AI y Artifact Registry, borra los recursos de Google Cloud que creaste durante este instructivo:

Para anular la implementación del modelo en el extremo y borrarlo, ejecuta el siguiente comando en tu shell:

ENDPOINT_ID=$(gcloud ai endpoints list \
   --region=LOCATION \
   --filter=display_name=ENDPOINT_NAME \
   --format="value(name)")

DEPLOYED_MODEL_ID=$(gcloud ai endpoints describe $ENDPOINT_ID \
   --region=LOCATION \
   --format="value(deployedModels.id)")

gcloud ai endpoints undeploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --deployed-model-id=$DEPLOYED_MODEL_ID

gcloud ai endpoints delete $ENDPOINT_ID \
   --region=LOCATION \
   --quiet

Reemplaza LOCATION por la región en la que creaste tu modelo en una sección anterior.

Para borrar tu modelo, ejecuta el siguiente comando en tu shell:

MODEL_ID=$(gcloud ai models list \
   --region=LOCATION \
   --filter=display_name=DEPLOYED_MODEL_NAME \
   --format="value(name)")

gcloud ai models delete $MODEL_ID \
   --region=LOCATION \
   --quiet

Reemplaza LOCATION por la región en la que creaste tu modelo en una sección anterior.

Para borrar tu repositorio de Artifact Registry y la imagen del contenedor que contiene, ejecuta el siguiente comando en tu shell:
```
gcloud artifacts repositories delete saxml \
  --location=LOCATION \
  --quiet
```
Reemplaza LOCATION por la región en la que creaste tu repositorio de Artifact Registry en una sección anterior.

Limitaciones

En Vertex AI Prediction, las Cloud TPU solo se admiten en us-west1. Para obtener más información, consulta Ubicaciones.

¿Qué sigue?

Aprende a implementar otros modelos de Saxml, como Llama2 y GPT-J.

Entrega modelos abiertos de Gemma con TPU en Vertex AI Prediction con Saxml

Formación

Gemma

Saxml

TPU

Antes de comenzar

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 el bucket de Cloud Storage

Copia el modelo al bucket de Cloud Storage

Crea un repositorio de Artifact Registry

Envía la imagen del contenedor a Artifact Registry

Implementa el modelo

Sube un modelo

Gemma 2B-it

Gemma 7B-it

Crea un extremo

Implementa el modelo en el extremo

Obtén predicciones en línea del modelo implementado

Realice una limpieza

Limitaciones

¿Qué sigue?