Diffuser des modèles ouverts Gemma à l'aide de TPU sur Vertex AI Prediction avec Saxml

Ce guide explique comment diffuser un grand modèle de langage (LLM) Gemma ouvert à l'aide de TPU (Tensor Processing Units) sur Vertex AI Prediction avec Saxml. Dans ce guide, vous allez télécharger sur Cloud Storage les modèles Gemma adaptés aux instructions 2B et 7B, puis les déployer sur Vertex AI Prediction, qui exécute Saxml sur des TPU.

Expérience

En diffusant Gemma à l'aide de TPU sur Vertex AI Prediction avec Saxml. Vous pouvez utiliser une solution d'IA gérée qui prend en charge l'infrastructure de bas niveau et offre un moyen économique de diffuser des LLM. Cette section décrit les principales technologies utilisées dans ce tutoriel.

Gemma

Gemma est un ensemble de modèles d'intelligence artificielle (IA) générative, légers et disponibles publiquement, publiés sous licence ouverte. Ces modèles d'IA sont disponibles pour exécution dans vos applications, votre matériel, vos appareils mobiles ou vos services hébergés. Vous pouvez utiliser les modèles Gemma pour la génération de texte, mais vous pouvez également les ajuster pour des tâches spécialisées.

Pour en savoir plus, consultez la documentation Gemma.

Saxml

Saxml est un système expérimental qui diffuse des modèles Paxml, JAX et PyTorch pour l'inférence. Dans le cadre de ce tutoriel, nous verrons comment diffuser Gemma sur des TPU plus rentables pour Saxml. La configuration des GPU est similaire. Saxml propose des scripts permettant de créer des conteneurs pour Vertex AI Prediction que nous allons utiliser dans ce tutoriel.

TPU

Les TPU sont des circuits intégrés propres aux applications (ASIC) développés spécifiquement par Google et permettant d'accélérer les frameworks de traitement de données tels que TensorFlow, PyTorch et JAX.

Ce tutoriel diffuse les modèles Gemma 2B et Gemma 7B. Vertex AI Prediction héberge ces modèles sur les pools de nœuds TPU v5e à hôte unique suivants:

Gemma 2B: hébergé dans un pool de nœuds TPU v5e avec une topologie 1x1 représentant une puce TPU. Le type de machine pour les nœuds est ct5lp-hightpu-1t.
Gemma 7B: hébergé dans un pool de nœuds TPU v5e avec une topologie 2x2 représentant quatre puces TPU. Le type de machine pour les nœuds est ct5lp-hightpu-4t.

Avant de commencer

Connectez-vous à votre compte Google Cloud. Si vous débutez sur Google Cloud, créez un compte pour évaluer les performances de nos produits en conditions réelles. Les nouveaux clients bénéficient également de 300 $ de crédits gratuits pour exécuter, tester et déployer des charges de travail.

Dans Google Cloud Console, sur la page de sélection du projet, sélectionnez ou créez un projet Google Cloud.

Accéder au sélecteur de projet

Vérifiez que la facturation est activée pour votre projet Google Cloud.

Activer les API Vertex AI API and Artifact Registry API.

Activer les API

Dans Google Cloud Console, sur la page de sélection du projet, sélectionnez ou créez un projet Google Cloud.

Accéder au sélecteur de projet

Vérifiez que la facturation est activée pour votre projet Google Cloud.

Activer les API Vertex AI API and Artifact Registry API.

Activer les API

Dans la console Google Cloud, activez Cloud Shell.

Activer Cloud Shell

En bas de la fenêtre de la console Google Cloud, une session Cloud Shell démarre et affiche une invite de ligne de commande. Cloud Shell est un environnement shell dans lequel Google Cloud CLI est déjà installé, et dans lequel des valeurs sont déjà définies pour votre projet actuel. L'initialisation de la session peut prendre quelques secondes.

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 :

Installez Google Cloud CLI.
Pour initialiser gcloudCLI, exécutez la commande suivante :
```
gcloud init
```

Consultez la documentation de Artifact Registry pour installer Docker.
Assurez-vous de disposer de quotas suffisants pour cinq puces TPU v5e pour Vertex AI Prediction.
Créez un compte Kaggle, si vous n'en possédez pas déjà un.

Accéder au modèle

Pour accéder aux modèles Gemma en vue du déploiement sur Vertex AI Prediction, vous devez vous connecter à la plate-forme Kaggle, signer le contrat de consentement de la licence et obtenir un jeton d'API Kaggle. Dans ce tutoriel, vous utilisez un secret Kubernetes pour les identifiants Kaggle.

Vous devez signer le contrat de consentement pour utiliser Gemma. Procédez comme suit :

Accédez à la page d'autorisation du modèle sur Kaggle.com.
Connectez-vous à Kaggle si ce n'est pas déjà fait.
Cliquez sur Demande d'accès.
Dans la section Choose Account for Consent (Choisir un compte pour le consentement), sélectionnez Verify via Kaggle Account (Valider via un compte Kaggle) pour utiliser votre compte Kaggle pour le consentement.
Acceptez les Conditions d'utilisation du modèle.

Générer un jeton d'accès

Pour accéder au modèle via Kaggle, vous avez besoin d'un jeton d'API Kaggle.

Pour générer un nouveau jeton si vous n'en possédez pas, procédez comme suit:

Dans votre navigateur, accédez aux paramètres Kaggle.
Dans la section API, cliquez sur Create New Token (Créer un jeton).

Un fichier nommé kaggle.json est téléchargé.

Importer le jeton d'accès dans Cloud Shell

Dans Cloud Shell, vous pouvez importer le jeton d'API Kaggle dans votre projet Google Cloud:

Dans Cloud Shell, cliquez sur Plus > Importer.
Sélectionnez "Fichier", puis cliquez sur Sélectionner des fichiers.
Ouvrez le fichier kaggle.json.
Cliquez sur Importer.

Créer le bucket Cloud Storage

Créer un bucket Cloud Storage pour stocker les points de contrôle du modèle.

Dans Cloud Shell, exécutez la commande ci-dessous.

gcloud storage buckets create gs://CHECKPOINTS_BUCKET_NAME

Remplacez CHECKPOINTS_BUCKET_NAME par le nom du bucket Cloud Storage qui stocke les points de contrôle du modèle.

Copier le modèle dans un bucket Cloud Storage

Dans Cloud Shell, exécutez la commande ci-dessous.

pip install kaggle --break-system-packages

# For Gemma 2B
mkdir -p /data/gemma_2b-it
kaggle models instances versions download google/gemma/pax/2b-it/1 --untar -p /data/gemma_2b-it
gsutil -m cp -R /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/

# For Gemma 7B
mkdir -p /data/gemma_7b-it
kaggle models instances versions download google/gemma/pax/7b-it/1 --untar -p /data/gemma_7b-it
gsutil -m cp -R /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/

Créer un dépôt Artifact Registry

Créez un dépôt Artifact Registry pour stocker l'image de conteneur que vous allez créer dans la section suivante.

Activez le service de l'API Artifact Registry pour votre projet.

gcloud services enable artifactregistry.googleapis.com

Exécutez la commande suivante dans votre shell pour créer le dépôt Artifact Registry :

gcloud artifacts repositories create saxml \
 --repository-format=docker \
 --location=LOCATION \
 --description="Saxml Docker repository"

Remplacez LOCATION par la région où Artifact Registry stocke votre image de conteneur. Vous devez ensuite créer une ressource de modèle Vertex AI sur un point de terminaison régional correspondant à cette région. Par conséquent, choisissez une région dans laquelle Vertex AI dispose d'un point de terminaison régional. Par exemple, us-west1 pour les TPU.

Transférer l'image de conteneur vers Artifact Registry

Le conteneur Saxml prédéfini est disponible sur us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest. Copiez-le dans Artifact Registry. Configurer Docker pour accéder à Artifact Registry Transférez ensuite votre image de conteneur vers votre dépôt Artifact Registry.

Pour autoriser votre installation Docker locale à transférer vers Artifact Registry dans la région choisie, exécutez la commande suivante dans votre interface système :
```
gcloud auth configure-docker LOCATION-docker.pkg.dev
```
- Remplacez LOCATION par la région dans laquelle vous avez créé votre dépôt.
Pour copier l'image de conteneur que vous venez de créer dans Artifact Registry, exécutez la commande suivante dans votre shell :
```
docker tag us-docker.pkg.dev/vertex-ai/prediction/sax-tpu:latest LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest
```
Pour transférer l'image de conteneur que vous venez de créer vers Artifact Registry, exécutez la commande suivante dans votre shell :
```
docker push LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest
```
Remplacez les éléments suivants, comme vous l'avez fait dans la section précédente :
- LOCATION : région de votre dépôt Artifact Registry.
- PROJECT_ID : ID de votre projet Google Cloud

Déployer le modèle

Importer un modèle

Pour importer une ressource Model utilisant votre conteneur Saxml, exécutez la commande gcloud ai models upload suivante:

Gemma 2B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma2BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Gemma 7B-it

gcloud ai models upload \
  --region=LOCATION \
  --display-name=DEPLOYED_MODEL_NAME \
  --container-image-uri=LOCATION-docker.pkg.dev/PROJECT_ID/saxml/saxml-tpu:latest \
  --artifact-uri='gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/' \
  --container-args='--model_path=saxml.server.pax.lm.params.gemma.Gemma7BFP16' \
  --container-args='--platform_chip=tpuv5e' \
  --container-args='--platform_topology=2x2' \
  --container-args='--ckpt_path_suffix=checkpoint_00000000' \
  --container-ports=8502

Remplacez les éléments suivants :

PROJECT_ID : ID de votre projet Google Cloud
LOCATION_ID : région dans laquelle vous utilisez Vertex AI. Notez que les TPU ne sont disponibles que dans la région us-west1.
DEPLOYED_MODEL_NAME : nom de l'élément DeployedModel. Vous pouvez également utiliser le nom à afficher du Model pour le DeployedModel.

Créer un point de terminaison

Vous devez déployer le modèle sur un point de terminaison avant de pouvoir l'utiliser pour diffuser des prédictions en ligne. Si vous déployez un modèle sur un point de terminaison existant, vous pouvez ignorer cette étape. L'exemple suivant utilise la commande gcloud ai endpoints create :

gcloud ai endpoints create \
  --region=LOCATION \
  --display-name=ENDPOINT_NAME

Remplacez les éléments suivants :

LOCATION_ID : région dans laquelle vous utilisez l'IA Vertex.
ENDPOINT_NAME : nom du point de terminaison à afficher.

La création du point de terminaison par l'outil Google Cloud CLI peut prendre quelques secondes.

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.

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

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

gcloud ai endpoints deploy-model $ENDPOINT_ID \
  --region=LOCATION \
  --model=$MODEL_ID \
  --display-name=DEPLOYED_MODEL_NAME \
  --machine-type=ct5lp-hightpu-4t \
  --traffic-split=0=100

Remplacez les éléments suivants :

LOCATION_ID : région dans laquelle vous utilisez l'IA Vertex.
ENDPOINT_NAME : nom à afficher du point de terminaison.
DEPLOYED_MODEL_NAME : nom de l'élément DeployedModel. Vous pouvez également utiliser le nom à afficher du Model pour le DeployedModel.

Gemma 2B peut être déployé sur une machine ct5lp-hightpu-1t plus petite. Dans ce cas, vous devez spécifier --platform_topology=1x1 lors de l'importation du modèle.

L'outil Google Cloud CLI peut prendre quelques secondes pour déployer le modèle sur le point de terminaison. Une fois le modèle déployé, cette commande affiche le résultat suivant :

  Deployed a model to the endpoint xxxxx. Id of the deployed model: xxxxx.

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

Pour appeler le modèle via le point de terminaison Vertex AI Prediction, formatez la requête de prédiction à l'aide d'un objet JSON de requête d'inférence standard.

L'exemple suivant utilise la commande gcloud ai endpoints predict :

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

gcloud ai endpoints predict $ENDPOINT_ID \
  --region=LOCATION \
  --http-headers=Content-Type=application/json \
  --json-request instances.json

Remplacez l'élément suivant :

LOCATION_ID : région dans laquelle vous utilisez l'IA Vertex.
ENDPOINT_NAME : nom du point de terminaison à afficher.
instances.json a le format suivant: {"instances": [{"text_batch": "<your prompt>"},{...}]}.

Effectuer un nettoyage

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

Pour annuler le déploiement du modèle sur le point de terminaison et supprimer le point de terminaison, exécutez la commande suivante dans votre shell :

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

Remplacez LOCATION par la région dans laquelle vous avez créé votre modèle dans une section précédente.

Pour supprimer votre modèle, exécutez la commande suivante dans votre interface système :

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

Remplacez LOCATION par la région dans laquelle vous avez créé votre modèle dans une section précédente.

Pour supprimer le dépôt Artifact Registry et l'image de conteneur, exécutez la commande suivante dans votre interface système :
```
gcloud artifacts repositories delete saxml \
  --location=LOCATION \
  --quiet
```
Remplacez LOCATION par la région dans laquelle vous avez créé votre dépôt Artifact Registry dans une section précédente.

Limites

Sur Vertex AI Prediction, les Cloud TPU ne sont compatibles qu'avec us-west1. Pour plus d'informations, consultez la section Emplacements.

Étapes suivantes

Découvrez comment déployer d'autres modèles Saxml tels que Llama2 et GPT-J.

Diffuser des modèles ouverts Gemma à l'aide de TPU sur Vertex AI Prediction avec Saxml

Expérience

Gemma

Saxml

TPU

Avant de commencer

Accéder au modèle

Signer le contrat de consentement de la licence

Générer un jeton d'accès

Importer le jeton d'accès dans Cloud Shell

Créer le bucket Cloud Storage

Copier le modèle dans un bucket Cloud Storage

Créer un dépôt Artifact Registry

Transférer l'image de conteneur vers Artifact Registry

Déployer le modèle

Importer un modèle

Gemma 2B-it

Gemma 7B-it

Créer un point de terminaison

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

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

Effectuer un nettoyage

Limites

Étapes suivantes