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 estct5lp-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 estct5lp-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.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the Vertex AI API.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the Vertex AI API.
-
In the Google Cloud console, 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 de2x2
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.
Signer le contrat de consentement de la licence
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 > Importer. Plus
- 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 duModel
pour leDeployedModel
.
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 duModel
pour leDeployedModel
.
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.