Se usó la API de Cloud Translation para traducir esta página.

Entrega modelos abiertos de Gemma mediante 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 con instrucciones de parámetros 2B y 7B en Cloud Storage y, luego, los implementarás en Vertex AI Prediction, que ejecuta Saxml en TPUs.

Fondo

Entregando Gemma con TPU en Vertex AI Prediction con Saxml Puedes aprovechar una solución de IA administrada que se encargue de la infraestructura de bajo nivel y ofrezca una forma 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 en el caso de tareas especializadas.

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

Saxml

Saxml es un sistema experimental que admite modelos Paxml, JAX y PyTorch para la inferencia. Solo para este instructivo, veremos cómo implementar Gemma en TPU que son más rentables para Saxml. La configuración de las GPUs es similar. Saxml ofrece secuencias de comandos para compilar contenedores para Vertex AI Prediction que usaremos en este instructivo.

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.

En este instructivo, se implementan 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 una 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 una topología 2x2 que representa cuatro chips TPU. El tipo de máquina para los nodos es ct5lp-hightpu-4t.

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.

Go to project selector

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

Enable the Vertex AI 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

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

Enable the Vertex AI API.

Enable the API

In the Google Cloud console, activate Cloud Shell.

Activate Cloud Shell

At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

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:

Install the Google Cloud CLI.
To initialize the gcloud CLI, run the following command:
```
gcloud init
```
Asegúrate de tener suficiente cuota para los chips TPU v5e en Vertex AI Prediction. De forma predeterminada, esta cuota es 0. Para una topología 1x1, debe ser 1. Para 2x2, debe ser 4. Para ejecutar ambas topologías, debe ser 5.
Crea una cuenta de Kaggle, si todavía no tienes una.

Obtén acceso al modelo

Ten en cuenta que es posible que Cloud Shell no tenga recursos suficientes para descargar los pesos del modelo. Si es así, puedes crear una instancia de Vertex AI Workbench para realizar esa tarea.

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

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.
Elige 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
gcloud storage cp /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/ --recursive

# 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
gcloud storage cp /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/ --recursive

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=us-docker.pkg.dev/vertex-ai/prediction/sax-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=us-docker.pkg.dev/vertex-ai/prediction/sax-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 en el extremo de Vertex AI Prediction, formatea la solicitud de predicción mediante 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>"},{...}]}

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

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 mediante TPU en Vertex AI Prediction con Saxml.

Fondo

Gemma

Saxml

TPU

Antes de comenzar

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

Copia el modelo al bucket de Cloud Storage

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

Realiza una limpieza

Limitaciones

¿Qué sigue?