Entrega modelos abiertos de Llama 3 con Cloud TPU de varios hosts en Gemini Enterprise Agent Platform con Saxml

Llama 3 es un modelo de lenguaje grande de código abierto (LLM) de Meta. En esta guía, se muestra cómo entregar un LLM de Llama 3 con unidades de procesamiento tensorial (TPU) de varios hosts en la plataforma de agentes de Gemini Enterprise con Saxml.

En esta guía, descargarás las ponderaciones y el tokenizador del modelo Llama 3 70B y los desplegarás en Gemini Enterprise Agent Platform, que ejecuta Saxml en TPUs.

Antes de comenzar

Te recomendamos que uses una VM M2 con optimización de memoria para descargar el modelo y convertirlo a Saxml. Esto se debe a que el proceso de conversión del modelo requiere una memoria significativa y puede fallar si eliges un tipo de máquina que no tiene suficiente memoria.

  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 and Artifact Registry APIs, if any are 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 APIs

  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.

  5. Sigue la documentación de Artifact Registry para instalar Docker.
  6. Asegúrate de tener cuotas suficientes para 16 chips TPU v5e en Gemini Enterprise Agent Platform.

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

Si utilizas un shell diferente en lugar de Cloud Shell para el despliegue del modelo asegúrese de que la versión de Google Cloud CLI es posterior a 475.0.0. Puedes actualizar la CLI de Google Cloud ejecutando el comando gcloud components update.

Si implementas tu modelo con el SDK de Vertex AI, asegúrate de tener la versión 1.50.0 o una posterior.

Obtén acceso al modelo y descarga los pesos del modelo

Los siguientes pasos son para una instancia de Agent Platform Workbench que tiene una VM con optimización de memoria M2. Si quieres obtener información para cambiar el tipo de máquina de una instancia de Agent Platform Workbench, consulta Cambia el tipo de máquina de una instancia de Agent Platform Workbench.

  1. Ve a la página de consentimiento del modelo de Llama.

  2. Selecciona Llama 3, completa el formulario de consentimiento y acepta los términos y condiciones.

  3. Busca en tu carpeta de entradaun correo electrónico con una URL firmada.

  4. Descarga la secuencia de comandos download.sh de GitHub ejecutando el siguiente comando:

    wget https://raw.githubusercontent.com/meta-llama/llama3/main/download.sh
chmod +x download.sh

  5. Para descargar los pesos del modelo, ejecuta el script download.sh que descargaste de GitHub.

  6. Cuando se te solicite, ingresa la URL firmada del correo electrónico que recibiste en la sección anterior.

  7. Cuando se te soliciten los modelos que desea descargar, introduce 70B.

Convierte las ponderaciones del modelo en el formato Saxml

  1. Ejecuta el siguiente comando para descargar Saxml:

    git clone https://github.com/google/saxml.git

  2. Ejecuta los siguientes comandos para configurar un entorno virtual de Python:

    python -m venv .
source bin/activate

  3. Ejecuta los siguientes comandos para instalar las dependencias:

    pip install --upgrade pip

pip install paxml

pip install praxis

pip install torch

  4. Para convertir las ponderaciones del modelo al formato Saxml, ejecuta el siguiente comando:

    python3 saxml/saxml/tools/convert_llama_ckpt.py \
    --base PATH_TO_META_LLAMA3 \
    --pax PATH_TO_PAX_LLAMA3 \
    --model-size llama3_70b

    Reemplaza lo siguiente:

    • PATH_TO_META_LLAMA3: la ruta al directorio que contiene las ponderaciones del modelo descargado
    • PATH_TO_PAX_LLAMA3: la ruta del directorio en el que almacenar los pesos del modelo convertido

    Los modelos convertidos se ubicarán en la carpeta $PATH_TO_PAX_LLAMA3/checkpoint_00000000.

  5. Copia el archivo tokenizer del directorio original en una subcarpeta denominada vocabs como se indica a continuación:

    cp $PATH_TO_META_LLAMA3/tokenizer.model $PATH_TO_PAX_LLAMA3/vocabs/tokenizer.model

  6. Añada un archivo vacío commit_success.txt en la carpeta $PATH_TO_PAX_LLAMA3 y en las subcarpetas metadata y state de dicha carpeta como sigue:

    touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/commit_success.txt
touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/commit_success.txt
touch $PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/commit_success.txt

    La carpeta $PATH_TO_PAX_LLAMA3 ahora contiene las siguientes carpetas y archivos:

    $PATH_TO_PAX_LLAMA3/checkpoint_00000000/commit_success.txt
$PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/
$PATH_TO_PAX_LLAMA3/checkpoint_00000000/metadata/commit_success.txt
$PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/
$PATH_TO_PAX_LLAMA3/checkpoint_00000000/state/commit_success.txt
$PATH_TO_PAX_LLAMA3/vocabs/tokenizer.model

Crea un bucket de Cloud Storage

Crea un bucket de Cloud Storage para almacenar las ponderaciones del modelo convertido.

  1. En Cloud Shell, ejecuta los siguientes comandos y reemplaza PROJECT_ID por el ID del proyecto:

    projectid=PROJECT_ID
gcloud config set project ${projectid}

  2. Para crear el bucket, ejecuta el siguiente comando: 

    gcloud storage buckets create gs://WEIGHTS_BUCKET_NAME

    Reemplaza WEIGHTS_BUCKET_NAME por el nombre que deseas usar para el bucket.

Copia las ponderaciones del modelo en el bucket de Cloud Storage

Para copiar las ponderaciones del modelo en tu bucket, ejecuta el siguiente comando:

gcloud storage cp PATH_TO_PAX_LLAMA3/* gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b/ --recursive

Sube el modelo

Hay disponible un contenedor Saxml precompilado en us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest.

Para subir un recurso Model a Gemini Enterprise Agent Platform con el contenedor de Saxml compilado previamente, ejecuta el comando gcloud ai models upload de la siguiente manera:

gcloud ai models upload \
    --region=LOCATION \
    --display-name=MODEL_DISPLAY_NAME \
    --container-image-uri=us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest \
    --artifact-uri='gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b' \
    --container-args='--model_path=saxml.server.pax.lm.params.lm_cloud.LLaMA3_70BFP16x16' \
    --container-args='--platform_chip=tpuv5e' \
    --container-args='--platform_topology=4x4' \
    --container-args='--ckpt_path_suffix=checkpoint_00000000' \
    --container-deployment-timeout-seconds=2700 \
    --container-ports=8502 \
    --project=PROJECT_ID

Realiza los siguientes reemplazos:

  • LOCATION: Es la región en la que usas Gemini Enterprise Agent Platform. Ten en cuenta que las TPU solo están disponibles en us-west1.
  • MODEL_DISPLAY_NAME: el nombre de visualización que deseas para tu modelo
  • PROJECT_ID: Es el ID de tu proyecto de Cloud de Confiance by S3NS .

Crea un extremo de inferencia en línea

Para crear el extremo, ejecuta el siguiente comando: 

gcloud ai endpoints create \
    --region=LOCATION \
    --display-name=ENDPOINT_DISPLAY_NAME \
    --project=PROJECT_ID

Reemplaza ENDPOINT_DISPLAY_NAME por el nombre visible que deseas para tu extremo.

Implementa el modelo con el extremo

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

En este tutorial, se despliega un modelo Llama 3 70B que se fragmenta para 16 Cloud TPU v5e utilizando una topología 4x4. Sin embargo, puedes especificar cualquiera de las siguientes topologías de Cloud TPU multihost compatibles:

Tipo de máquina Topología Cantidad de chips TPU Cantidad de hosts
ct5lp-hightpu-4t 4x4 16 2
ct5lp-hightpu-4t 4x8 32 4
ct5lp-hightpu-4t 8x8 64 8
ct5lp-hightpu-4t 8x16 128 16
ct5lp-hightpu-4t 16x16 256 32

Si implementas un modelo de Llama diferente que se define en el repositorio de GitHub de Saxml, asegúrate de que esté particionado para que coincida con la cantidad de dispositivos a los que te orientas y que Cloud TPU tiene suficiente memoria para cargar el modelo.

Para obtener más información sobre la implementación de un modelo en Cloud TPU de host único, consulta Implementa un modelo.

Para obtener una lista completa de los tipos y las regiones de Cloud TPU compatibles, consulta Ubicaciones de Agent Platform de Gemini Enterprise.

  1. Obtén el ID del extremo de inferencia en línea:

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

  2. Obtén el ID de tu modelo:

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

  3. Implementa el modelo en el extremo:

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

    Reemplaza DEPLOYED_MODEL_NAME por un nombre para el elemento implementado. Puede ser el mismo que el nombre visible del modelo (MODEL_DISPLAY_NAME).

    Es posible que se agote el tiempo de espera de la operación de implementación.

    El comando deploy-model devuelve un ID de operación que puede utilizarse para comprobar cuándo ha finalizado la operación. Puedes sondear el estado de la operación hasta que la respuesta incluya "done": true. Utiliza el siguiente comando para sondear el estado:

    gcloud ai operations describe \
--region=LOCATION \
OPERATION_ID

    Sustituye OPERATION_ID por el ID de operación que mostró el comando anterior.

Obtén inferencias en línea del modelo implementado

Para obtener inferencias en línea del extremo de Gemini Enterprise Agent Platform, ejecuta el comando gcloud ai endpoints predict.

  1. Ejecuta el siguiente comando para crear un archivo request.json que contenga una solicitud de inferencia de muestra:

    cat << EOF > request.json
{"instances": [{"text_batch": "the distance between Earth and Moon is "}]}
EOF

  2. Para enviar la solicitud de inferencia en línea al extremo, ejecuta el siguiente comando:

    gcloud ai endpoints predict $ENDPOINT_ID \
    --project=PROJECT_ID \
    --region=LOCATION \
    --json-request=request.json

Realiza una limpieza

Para evitar incurrir en más cargos de Gemini Enterprise Agent Platform, 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 los siguientes comandos:

    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

  2. Para borrar tu modelo, ejecuta los siguientes comandos:

    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