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

Vista previa

Esta función está sujeta a las "Condiciones de las Ofertas de Fase Previa a la DG" de la sección Condiciones Generales del Servicio de las Condiciones Específicas del Servicio. Las funciones de fase previa a la DG están disponibles “tal cual” y podrían tener asistencia limitada. Para obtener más información, consulta las descripciones de la etapa de lanzamiento.

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

Fondo

Entregando Gemma mediante 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 los 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 GPU es similar. Saxml ofrece secuencias de comandos para compilar contenedores de Vertex AI Prediction que usaremos en este instructivo.

TPU

Las TPU son circuitos integrados personalizados específicos para aplicaciones de Google (ASIC) que se usan para acelerar 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

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

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

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

   Go to project selector

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

4. Enable the Vertex AI API.

   Enable the API

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

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

   Go to project selector

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

7. Enable the Vertex AI API.

   Enable the API

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

1. Install the Google Cloud CLI.

2. To initialize the gcloud CLI, run the following command:

   gcloud init

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

Firma el acuerdo de consentimiento de licencia

Debes firmar el acuerdo de consentimiento para usar Gemma. Sigue estas instrucciones:

1. Accede a la página de consentimiento del modelo en Kaggle.com.
2. Ingresa en Kaggle si aún no lo has hecho.
3. Haz clic en Solicitar acceso.
4. En la sección Elegir cuenta para el consentimiento, selecciona Verificar mediante la cuenta de Kaggle para usar tu cuenta de Kaggle para obtener el consentimiento.
5. Acepta los Términos y Condiciones del modelo.

Genera un token de acceso

Para acceder al modelo a través de Kaggle, necesitas un token de la API de Kaggle.

Sigue estos pasos para generar un token nuevo si aún no tienes uno:

1. En el navegador, ve a la configuración de Kaggle.

2. En la sección API, haz clic en Crear token nuevo.

   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:

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

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

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

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