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.

Contexte

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

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.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the Vertex AI API.

Enable the API

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Go to project selector

Make sure that billing is enabled for your Google Cloud project.

Enable the Vertex AI API.

Enable the API

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.

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 :

Install the Google Cloud CLI.
To initialize the gcloud CLI, run the following command:
```
gcloud init
```
Assurez-vous de disposer d'un quota suffisant pour les puces TPU v5e pour Vertex AI Prediction. Par défaut, ce quota est défini sur 0. Pour une topologie 1x1, la valeur doit être 1. La valeur de 2x2 doit être 4. Pour exécuter les deux topologies, vous devez indiquer 5.
Créez un compte Kaggle, si vous n'en possédez pas déjà un.

Accéder au modèle

Notez que Cloud Shell peut ne pas disposer de ressources suffisantes pour télécharger les pondérations de modèle. Si tel est le cas, vous pouvez créer une instance Vertex AI Workbench pour effectuer cette tâche.

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
gcloud storage cp /data/gemma_2b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/ --recursive

# 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
gcloud storage cp /data/gemma_7b-it/* gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/ --recursive

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=us-docker.pkg.dev/vertex-ai/prediction/sax-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=us-docker.pkg.dev/vertex-ai/prediction/sax-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>"},{...}]}.

Nettoyer

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.

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

Contexte

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

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é

Nettoyer

Limites

Étapes suivantes