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

Preview

Cette fonctionnalité est soumise(s) aux "Conditions des offres de pré-DG" de la section "Conditions générales des Services" des Conditions spécifiques du service. Les fonctionnalités de pré-DG sont disponibles "en l'état" et peuvent avoir une prise en charge 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 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 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.

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.

In the Google Cloud console, on the project selector page, select or create a Google Cloud 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 (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

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

Verify that billing is enabled for your Google Cloud project.

Enable the Vertex AI and Artifact Registry APIs.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the APIs

In the Google Cloud console, on the project selector page, select or create a Google Cloud 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 (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

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

Verify that billing is enabled for your Google Cloud project.

Enable the Vertex AI and Artifact Registry APIs.

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

Enable the APIs

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

2. Consultez la documentation de Artifact Registry pour installer Docker.
3. Assurez-vous de disposer de quotas suffisants pour 16 puces TPU v5e pour Vertex AI.

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. If you're using an external identity provider (IdP), you must first sign in to the gcloud CLI with your federated identity.

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

   gcloud init

Si vous utilisez un autre shell que Cloud Shell pour le déploiement de modèles, 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 d'autorisation du modèle Llama.

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

3. Consultez votre boîte de réception pour trouver l'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 figurant dans l'e-mail que vous avez reçu dans la section précédente.

7. Lorsque vous êtes invité à indiquer les modèles à télécharger, 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 seront placés dans le dossier $PATH_TO_PAX_LLAMA3/checkpoint_00000000.

5. Copiez le fichier du tokenizer depuis le 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, ainsi que les sous-dossiers metadata et state dans 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 à 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 d'inférence 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 que vous souhaitez attribuer à 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 Cloud TPU multi-hôtes 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é pour correspondre au nombre d'appareils que vous ciblez et que Cloud TPU dispose de suffisamment de mémoire pour charger le modèle.

Pour savoir comment déployer un modèle sur des Cloud TPU à hôte unique, consultez Déployer un modèle.

Pour obtenir la liste complète des types et régions Cloud TPU compatibles, consultez la page Emplacements Vertex AI.

1. Obtenez l'ID du point de terminaison d'inférence 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 s'agir du même nom que celui 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. Exécutez 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 inférences en ligne à partir du modèle déployé

Pour obtenir des inférences en ligne à partir du point de terminaison Vertex AI, 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 d'inférence :

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

2. Pour envoyer la requête d'inférence 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 Google Cloud ressources 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