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

    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.

    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 Llamar 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. Ejecuta el siguiente comando para descargar la secuencia de comandos download.sh desde GitHub:

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

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 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, 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 de la 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 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 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