Diffuser des modèles ouverts Gemma à l'aide de TPU sur Gemini Enterprise Agent Platform avec Saxml

Ce guide explique comment diffuser un grand modèle de langage (LLM) Gemma ouvert à l'aide de TPU (Tensor Processing Units) sur la plate-forme Gemini Enterprise Agent 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 Gemini Enterprise Agent Platform, qui exécute Saxml sur des TPU.

Arrière-plan

En diffusant Gemma à l'aide de TPU sur Gemini Enterprise Agent Platform 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 Gemini Enterprise Agent Platform 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. Gemini Enterprise Agent Platform 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

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

    Go to project selector

  2. Verify that billing is enabled for your Cloud de Confiance project.

  3. Enable the Agent Platform API, if it is not already enabled.

    Roles required to enable APIs

    To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

    Enable the API

  4. Dans la console Cloud de Confiance , activez Cloud Shell.

    Activer Cloud Shell

    En bas de la console Cloud de Confiance , 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 Cloud de Confiance by S3NS. Si vous souhaitez utiliser un autre shell au lieu de Cloud Shell, exécutez la configuration supplémentaire suivante :

  1. Installez la Google Cloud CLI.

  2. Configurez la gcloud CLI afin d'utiliser votre identité fédérée.

    Pour en savoir plus, consultez Se connecter à la gcloud CLI avec votre identité fédérée.

  3. Pour initialiser la gcloud CLI, exécutez la commande suivante :

    gcloud init
  4. Assurez-vous de disposer d'un quota suffisant pour les puces TPU v5e pour Gemini Enterprise Agent Platform. 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.
  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 Gemini Enterprise Agent Platform Workbench pour effectuer cette tâche.

Pour accéder aux modèles Gemma en vue du déploiement sur Gemini Enterprise Agent Platform, 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 :

  1. Accédez à la page d'autorisation du modèle sur Kaggle.com.
  2. Connectez-vous à Kaggle si ce n'est pas déjà fait.
  3. Cliquez sur Demande d'accès.
  4. 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.
  5. 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:

  1. Dans votre navigateur, accédez aux paramètres Kaggle.
  2. 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 Cloud de Confiance by S3NS :

  1. Dans Cloud Shell, cliquez sur Plus > Importer.
  2. Sélectionnez "Fichier", puis cliquez sur Sélectionner des fichiers.
  3. Ouvrez le fichier kaggle.json.
  4. 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 projetCloud de Confiance
  • LOCATION_ID : région dans laquelle vous utilisez la plate-forme d'agents. 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 inférences 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 la plate-forme d'agents.
  • ENDPOINT_NAME : nom à afficher pour le point de terminaison.

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 la plate-forme d'agents.
  • ENDPOINT_NAME : nom à afficher pour le 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 inférences en ligne à partir du modèle déployé

Pour appeler le modèle via le point de terminaison Gemini Enterprise Agent Platform, formatez la requête d'inférence à 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 les éléments suivants :

  • LOCATION_ID : région dans laquelle vous utilisez la plate-forme d'agents.
  • ENDPOINT_NAME : nom à afficher pour le point de terminaison.
  • instances.json a le format suivant: {"instances": [{"text_batch": "<your prompt>"},{...}]}.

Nettoyer

Pour éviter que des frais liés à Gemini Enterprise Agent Platform et des frais liés à Artifact Registry supplémentaires vous soient facturés, supprimez les Cloud de Confiance 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 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.

  2. 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 Gemini Enterprise Agent Platform, 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.