Diffuser des modèles ouverts Llama 3 à l'aide de TPU multi-hôtes sur Vertex AI Prediction avec Saxml

Preview

Cette fonctionnalité est soumise aux "Conditions des offres de pré-DG" dans la section "Conditions d'utilisation générales" des Conditions spécifiques au service. Les fonctionnalités pré-DG sont disponibles "en l'état" et peuvent avoir une assistance limitée. Pour en savoir plus, consultez les descriptions des étapes de lancement.

Llama 3 est un grand modèle de langage (LLM) Open Source de Meta. Ce guide explique comment diffuser un LLM Llama 3 à l'aide de Tensor Processing Units (TPU) à hôtes multiples sur Vertex AI Prediction avec Saxml.

Dans ce guide, vous téléchargez les pondérations de modèle Llama 3 70B et la fonction de tokenisation, puis vous les déployez sur Vertex AI Prediction qui exécute Saxml sur des TPU.

Remarque : Une version GPU uniquement de Llama 3 est disponible dans Model Garden. Pour en savoir plus sur Model Garden, consultez la page Explorer les modèles d'IA dans Model Garden.

Avant de commencer

Nous vous recommandons d'utiliser une VM M2 à mémoire optimisée pour télécharger le modèle et le convertir au format Saxml. En effet, le processus de conversion du modèle nécessite une quantité importante de mémoire et peut échouer si vous choisissez un type de machine qui ne dispose pas de suffisamment de mémoire.

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. Consultez la documentation de Artifact Registry pour installer Docker.
10. Assurez-vous de disposer de quotas suffisants pour 16 puces TPU v5e pour Vertex AI Prediction.

Ce tutoriel suppose que vous utilisez Cloud Shell pour interagir avec Google Cloud. Si vous souhaitez utiliser un autre shell au lieu de Cloud Shell, exécutez la configuration supplémentaire suivante :

1. Install the Google Cloud CLI.

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

   gcloud init

Si vous utilisez une interface système différente au lieu de Cloud Shell pour le déploiement du modèle, assurez-vous que la version de Google Cloud CLI est ultérieure à 475.0.0. Vous pouvez mettre à jour Google Cloud CLI en exécutant la commande gcloud components update.

Si vous déployez votre modèle à l'aide du SDK Vertex AI, assurez-vous de disposer de la version 1.50.0 ou d'une version ultérieure.

Accéder au modèle et télécharger les poids du modèle

Les étapes suivantes concernent une instance Vertex AI Workbench dotée d'une VM M2 à mémoire optimisée. Pour savoir comment modifier le type de machine d'une instance Vertex AI Workbench, consultez la section Modifier le type de machine d'une instance Vertex AI Workbench.

1. Accédez à la page de consentement du modèle Llama.

2. Sélectionnez Llama 3, remplissez le formulaire de consentement et acceptez les conditions d'utilisation.

3. Vérifiez que vous avez bien reçu un e-mail contenant une URL signée.

4. Téléchargez le script download.sh à partir de GitHub en exécutant la commande suivante :

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

5. Pour télécharger les pondérations du modèle, exécutez le script download.sh que vous avez téléchargé depuis GitHub.

6. Lorsque vous y êtes invité, saisissez l'URL signée de l'e-mail que vous avez reçu dans la section précédente.

7. Lorsque vous êtes invité à télécharger les modèles, saisissez 70B.

Convertir les pondérations du modèle au format Saxml

1. Exécutez la commande suivante pour télécharger Saxml :

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

2. Exécutez les commandes suivantes pour configurer un environnement virtuel Python :

   python -m venv .
   source bin/activate
   

3. Exécutez les commandes suivantes pour installer les dépendances :

   pip install --upgrade pip
   
   pip install paxml
   
   pip install praxis
   
   pip install torch
   

4. Pour convertir les pondérations du modèle au format Saxml, exécutez la commande suivante :

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

   Remplacez les éléments suivants :

   • PATH_TO_META_LLAMA3 : chemin d'accès au répertoire contenant les pondérations du modèle téléchargé
   • PATH_TO_PAX_LLAMA3 : chemin d'accès au répertoire dans lequel stocker les pondérations du modèle converti
   Remarque : Vous pouvez utiliser cette commande pour n'importe quel modèle Llama 2 ou Llama 3.

   Les modèles convertis sont placés dans le dossier $PATH_TO_PAX_LLAMA3/checkpoint_00000000.

5. Copiez le fichier de générateur de jetons du répertoire d'origine dans un sous-dossier nommé vocabs comme suit :

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

6. Ajoutez un fichier commit_success.txt vide dans le dossier $PATH_TO_PAX_LLAMA3 et les sous-dossiers metadata et state de ce dossier, comme suit :

   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
   

   Le dossier $PATH_TO_PAX_LLAMA3 contient désormais les dossiers et fichiers suivants :

   $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
   

Créer un bucket Cloud Storage

Créer un bucket Cloud Storage pour stocker les pondérations du modèle converti.

1. Dans Cloud Shell, exécutez les commandes suivantes, en remplaçant PROJECT_ID par votre ID de projet :

   projectid=PROJECT_ID
   gcloud config set project ${projectid}
   

2. Pour créer le bucket, exécutez la commande suivante :

   gcloud storage buckets create gs://WEIGHTS_BUCKET_NAME
   

   Remplacez WEIGHTS_BUCKET_NAME par le nom que vous souhaitez utiliser pour le bucket.

Copier les pondérations du modèle dans le bucket Cloud Storage

Pour copier les pondérations du modèle dans votre bucket, exécutez la commande suivante :

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

Importer le modèle.

Le conteneur Saxml prédéfini est disponible sur us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest.

Pour importer une ressource Model dans Vertex AI Prediction à l'aide du conteneur Saxml prédéfini, exécutez la commande gcloud ai models upload comme suit :

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

Effectuez les remplacements suivants :

• LOCATION : région dans laquelle vous utilisez Vertex AI. Notez que les TPU ne sont disponibles que dans us-west1.
• MODEL_DISPLAY_NAME : nom à afficher souhaité pour votre modèle.
• PROJECT_ID : ID de votre projet Google Cloud

Remarque : Vous pouvez utiliser cette commande pour tout modèle Llama répertorié sur le lien GitHub suivant : https://github.com/google/saxml/blob/main/saxml/server/pax/lm/params/lm_cloud.py

Créer un point de terminaison de prédiction en ligne.

Pour créer le point de terminaison, exécutez la commande suivante :

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

Remplacez ENDPOINT_DISPLAY_NAME par le nom à afficher pour votre point de terminaison.

Déployer le modèle sur le point de terminaison

Une fois le point de terminaison prêt, déployez le modèle sur ce point de terminaison.

Dans ce tutoriel, vous allez déployer un modèle Llama 3 70B segmenté pour 16 puces Cloud TPU v5e à l'aide de la topologie 4x4. Toutefois, vous pouvez spécifier l'une des topologies multi-hôtes Cloud TPU compatibles suivantes :

Type de machine	Topologie	Nombre de puces TPU	Nombre d'hôtes
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 vous déployez un autre modèle Llama défini dans le dépôt GitHub Saxml, assurez-vous qu'il est partitionné en fonction du nombre d'appareils que vous ciblez et que Cloud TPU dispose de suffisamment de mémoire pour charger le modèle.

Pour en savoir plus sur le déploiement d'un modèle sur des Cloud TPU à hôte unique, consultez la page Déployer un modèle.

Pour en savoir plus sur les types Cloud TPU v5e, consultez TPU v5e.

1. Obtenez l'ID du point de terminaison de prédiction en ligne :

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

2. Obtenez l'ID de votre modèle :

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

3. Déployez le modèle sur le point de terminaison :

   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
   

   Remplacez DEPLOYED_MODEL_NAME par le nom du déploiement. Il peut être identique au nom à afficher du modèle (MODEL_DISPLAY_NAME).

   L'opération de déploiement peut expirer.

   La commande deploy-model renvoie un ID d'opération qui permet de vérifier la fin de l'opération. Vous pouvez interroger l'état de l'opération jusqu'à ce que la réponse indique "done": true. Utilisez la commande suivante pour interroger l'état :

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

   Remplacez OPERATION_ID par l'ID d'opération renvoyé par la commande précédente.

Obtenir des prédictions en ligne à partir du modèle déployé

Pour obtenir des prédictions en ligne à partir du point de terminaison Vertex AI Prediction, exécutez la commande gcloud ai endpoints predict.

1. Exécutez la commande suivante pour créer un fichier request.json contenant un exemple de requête de prédiction :

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

2. Pour envoyer la requête de prédiction en ligne au point de terminaison, exécutez la commande suivante :

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

Effectuer un nettoyage

Pour éviter que des frais liés à Vertex AI supplémentaires vous soient facturés, supprimez les ressources Google Cloud que vous avez créées au cours de ce tutoriel :

1. Pour annuler le déploiement du modèle sur le point de terminaison et supprimer le point de terminaison, exécutez les commandes suivantes :

   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. Pour supprimer votre modèle, exécutez les commandes suivantes :

   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