Entrega modelos abiertos de Gemma con TPUs en Gemini Enterprise Agent Platform 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 Gemini Enterprise Agent Platform 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 Gemini Enterprise Agent Platform, que ejecuta Saxml en TPUs.

Fondo

Entrega Gemma con TPU en Gemini Enterprise Agent Platform con Saxml. Puedes aprovechar una solución de IA administrada que se encarga de la infraestructura de bajo nivel y ofrece una forma rentable de entregar LLMs. 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 Gemini Enterprise Agent Platform 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. Gemini Enterprise Agent Platform 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. In the Cloud de Confiance console, on the project selector page, select or create a Cloud de Confiance project.

    Roles required to select or create a project

    • Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
    • Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

    Go to project selector

  2. Verify that billing is enabled for your Cloud de Confiance project.

  3. Enable the Agent Platform API, if it is not already enabled.

    Roles required to enable APIs

    To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

    Enable the API

  4. En la consola de Cloud de Confiance , activa Cloud Shell.

    Activa Cloud Shell

    En la parte inferior de la consola de Cloud de Confiance , 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 Cloud de Confiance by S3NS. Si quieres usar un shell diferente en lugar de Cloud Shell, realiza la siguiente configuración adicional:

  1. Instala Google Cloud CLI.

  2. Configura gcloud CLI para usar tu identidad federada.

    Para obtener más información, consulta Accede a la gcloud CLI con tu identidad federada.

  3. Para inicializar gcloud CLI, ejecuta el siguiente comando:

    gcloud init
  4. Asegúrate de tener suficiente cuota para los chips TPU v5e en Gemini Enterprise Agent Platform. 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.
  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 suficientes recursos para descargar los pesos del modelo. Si es así, puedes crear una instancia de Gemini Enterprise Agent Platform Workbench para realizar esa tarea.

Para obtener acceso a los modelos de Gemma para la implementación en Gemini Enterprise Agent Platform, 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:

  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 Cloud de Confiance by S3NS:

  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: Es el ID de tu Cloud de Confianceproyecto.
  • LOCATION_ID: Es la región en la que usas Agent Platform. 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.

Crear un extremo

Debes implementar el modelo en un extremo antes de que pueda usarse para entregar inferencias 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: Es la región en la que usas Agent Platform.
  • 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: Es la región en la que usas Agent Platform.
  • 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 inferencias en línea del modelo implementado

Para invocar el modelo a través del extremo de Gemini Enterprise Agent Platform, formatea la solicitud de inferencia 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: Es la región en la que usas Agent Platform.
  • 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 Gemini Enterprise Agent Platform y Artifact Registry, borra los Cloud de Confiance recursos 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 Gemini Enterprise Agent Platform, las Cloud TPU solo se admiten en us-west1. Para obtener más información, consulta Ubicaciones.

¿Qué sigue?