Entrega modelos abiertos de Llama 3 mediante Cloud TPU de varios hosts 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.

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.

Nota: Una versión de Llama 3 solo para GPU está disponible en Model Garden. Para obtener más información acerca de Model Garden, consulta Explora los modelos de IA en Model Garden.

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. 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 and Artifact Registry APIs.

   Enable the APIs

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 and Artifact Registry APIs.

   Enable the APIs

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.

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

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

   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
   Nota: Puedes utilizar este comando para cualquier modelo de Llama 2 o Llama 3.

   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:

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 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_ID: El ID de tu proyecto de Google Cloud.

Nota: Puedes usar este comando para cualquier modelo de Llama que aparezca en el siguiente vínculo de GitHub: https://github.com/google/saxml/blob/main/saxml/server/pax/lm/params/lm_cloud.py

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
   

Realiza una limpieza

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