Entrega modelos abiertos de Llama 3 mediante Cloud TPU de varios hosts en Vertex AI Prediction con Saxml

Llama 3 es un modelo de lenguaje grande de código abierto (LLM) de Meta. Esta guía le muestra cómo servir un LLM Llama 3 usando unidades de tratamiento de datos tensorial (TPUs) en Vertex AI Prediction con Saxml.

En esta guía, descargarás las ponderadaciones y el tokenizador del modelo Llama 3 70B y los desplegarás en Vertex AI Prediction, que ejecuta Saxml en TPUs.

Antes de comenzar

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

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

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

    Go to project selector

  3. Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.

  4. Habilita las API de Vertex AI and Artifact Registry.

    Habilita las API

    5.

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

    Go to project selector

  6. Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.

  7. Habilita las API de Vertex AI and Artifact Registry.

    Habilita las API

    8.

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

  9. Sigue la documentación de Artifact Registry para instalar Docker.
  10. Asegúrate de tener cuotas suficientes para 16 chips TPU v5e en Vertex AI Prediction.

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. Instala Google Cloud CLI.

  2. Para inicializar la CLI de gcloud, 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 Vertex AI Workbench que tiene una VM con optimización de memoria M1. Si quieres obtener información para cambiar el tipo de máquina de una instancia de Vertex AI Workbench, consulta Cambia el tipo de máquina de una instancia de Vertex AI 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 mediante la ejecución del 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 los siguientes archivos y carpetas:

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

gsutil -m cp -R PATH_TO_PAX_LLAMA3/* gs://WEIGHTS_BUCKET_NAME/llama3/pax_70b/

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 Vertex AI Prediction mediante el contenedor saxml compilado de forma previa, 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: la región en la que usas Vertex AI. Ten en cuenta que las TPU solo están disponibles en us-west1.
  • MODEL_DISPLAY_NAME: el nombre de visualización que desea para su modelo
  • PROJECT_IDEl ID de tu proyecto de Google Cloud.

Crea un extremo de predicción 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 el 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, puede especificar cualquiera de las siguientes topologías de TPU en nube 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 más información sobre los tipos de Cloud TPU v5e, consulta TPU v5e.

  1. Obtén el ID de extremo para el extremo de predicción 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 la implementación. Puede ser igual al 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 predicciones en línea del modelo implementado

Para obtener predicciones en línea del punto final Vertex AI Prediction ejecuta el comando gcloud ai endpoints predict.

  1. Ejecute el siguiente comando request.json para crear un archivo que contenga una ejemplo de solicitud de predicción:

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

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

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

Limpia

Para evitar incurrir en más cargos de Vertex AI, elimina los recursos de Google Cloud que creaste durante este tutorial:

  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