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


Ce guide explique comment diffuser un grand modèle de langage (LLM) Gemma à l'aide des TPU (Tensor Processing Units) sur Google Kubernetes Engine (GKE) avec Saxml. Dans ce guide, vous téléchargez les modèles Gemma adaptés aux instructions de paramètres 2B et 7B sur Cloud Storage, puis vous les déployez sur un cluster GKE Standard à l'aide de conteneurs exécutant Saxml.

Ce guide est un bon point de départ si vous avez besoin de l'évolutivité, de la résilience et de la rentabilité offertes par les fonctionnalités de Kubernetes lors du déploiement de votre modèle sur Saxml.

Si vous avez besoin d'une plate-forme d'IA gérée unifiée pour créer et diffuser rapidement des modèles de ML à moindre coût, nous vous recommandons d'essayer notre solution de déploiement Vertex AI.

Contexte

En diffusant Gemma à l'aide de TPU sur GKE avec Saxml, vous pouvez mettre en œuvre une solution de diffusion d'inférences robuste et prête pour la production avec tous les avantages de Kubernetes géré, y compris une évolutivité efficace et une meilleure disponibilité. Cette section décrit les principales technologies utilisées dans ce tutoriel.

Gemma

Gemma est un ensemble de modèles d'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.

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.

Avant d'utiliser des TPU dans GKE, nous vous recommandons de suivre le parcours de formation suivant :

  1. Découvrez la disponibilité actuelle des versions de TPU avec l'architecture système de Cloud TPU.
  2. Apprenez-en plus sur les TPU dans GKE.

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

  • Gemma 2B: modèle adapté aux instructions 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: modèle adapté aux instructions 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.

Saxml

Saxml est un système expérimental qui diffuse des modèles Paxml, JAX et PyTorch pour l'inférence. Le système Saxml comprend les composants suivants:

  • Cellule Saxml ou cluster Sax: serveur d'administration et groupe de serveurs de modèles. Le serveur d'administration assure le suivi des serveurs de modèles, attribue des modèles publiés aux serveurs de modèles à diffuser et aide les clients à localiser les serveurs de modèles diffusant des modèles publiés spécifiques.
  • Client Saxml: interface de programmation destinée à l'utilisateur pour le système Saxml. Le client Saxml inclut un outil de ligne de commande (saxutil) et une suite de bibliothèques clientes en Python, C++ et Go.

Dans ce tutoriel, vous utiliserez également le serveur HTTP Saxml. Le serveur HTTP Saxml est un serveur HTTP personnalisé qui encapsule la bibliothèque cliente Python Saxml et expose les API REST pour interagir avec le système Saxml. Les API REST incluent des points de terminaison permettant de publier, de répertorier, d'annuler la publication de modèles et de générer des prédictions.

Objectifs

Ce tutoriel est destiné aux clients d'IA générative qui utilisent JAX, ainsi qu'aux utilisateurs nouveaux ou existants de GKE qui souhaitent utiliser les fonctionnalités d'orchestration de conteneurs Kubernetes pour diffuser Gemma, tels que les ingénieurs en ML ou les ingénieurs MLOps (DevOps) et les administrateurs de plate-forme.

Ce tutoriel couvre les étapes suivantes :

  1. Préparez un cluster GKE standard avec la topologie TPU recommandée en fonction des caractéristiques du modèle.
  2. Déployer des composants Saxml sur GKE.
  3. Obtenez et publiez le modèle Gemma de paramètre 2B ou 7B.
  4. Diffuser les modèles publiés et interagir avec eux.

Architecture

Cette section décrit l'architecture GKE utilisée dans ce tutoriel. L'architecture comprend un cluster GKE Standard qui provisionne des TPU et héberge des composants Saxml pour déployer et diffuser des modèles Gemma 2B ou 7B. Le schéma suivant montre les composants de cette architecture:

Schéma de l'architecture déployée dans ce tutoriel

Cette architecture comprend les composants suivants :

  • Un cluster zonal GKE Standard.
  • Un pool de nœuds de tranche de TPU à hôte unique qui dépend du modèle Gemma que vous souhaitez diffuser :
    • Gemma 2B: configuré avec un TPU v5e avec une topologie 1x1. Une instance du serveur de modèles Saxml est configurée pour utiliser ce pool de nœuds.
    • Gemma 7B: configuré avec un TPU v5e avec une topologie 2x2. Une instance du serveur de modèles Saxml est configurée pour utiliser ce pool de nœuds.
  • Un pool de nœuds de processeur par défaut sur lequel le serveur d'administration Saxml et le serveur HTTP Saxml sont déployés.
  • Deux buckets Cloud Storage :
    • Un bucket Cloud Storage stocke l'état géré par un serveur d'administration.
    • Un bucket Cloud Storage stocke les points de contrôle du modèle Gemma.

Cette architecture présente les caractéristiques suivantes:

  • Un dépôt Artifact Registry public gère les images de conteneurs pour les composants Saxml.
  • Le cluster GKE utilise la fédération d'identité de charge de travail pour GKE. Tous les composants Saxml utilisent une fédération d'identité de charge de travail qui intègre un compte de service IAM pour accéder à des services externes tels que des buckets Cloud Storage.
  • Les journaux générés par les composants Saxml sont intégrés à Cloud Logging.
  • Vous pouvez utiliser Cloud Monitoring pour analyser les métriques de performances des pools de nœuds GKE, tels que ceux créés dans ce tutoriel.

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.
  • 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 required 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 required API.

    Enable the API

  • Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin

    Check for the roles

    1. In the Google Cloud console, go to the IAM page.

      Go to IAM
    2. Select the project.
    3. In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.

    4. For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.

    Grant the roles

    1. In the Google Cloud console, go to the IAM page.

      Accéder à IAM
    2. Sélectionnez le projet.
    3. Cliquez sur Accorder l'accès.
    4. Dans le champ Nouveaux comptes principaux, saisissez votre identifiant utilisateur. Il s'agit généralement de l'adresse e-mail d'un compte Google.

    5. Dans la liste Sélectionner un rôle, sélectionnez un rôle.
    6. Pour attribuer des rôles supplémentaires, cliquez sur Ajouter un autre rôle et ajoutez chaque rôle supplémentaire.
    7. Cliquez sur Enregistrer.
  • Assurez-vous de disposer de quotas suffisants pour cinq puces TPU v5e. Dans ce tutoriel, vous utilisez des instances à la demande.
  • Créez un compte Kaggle, si vous n'en possédez pas.

Préparer l'environnement pour Gemma

Lancer Cloud Shell

Dans ce tutoriel, vous utilisez Cloud Shell pour gérer les ressources hébergées sur Google Cloud. Cloud Shell est préinstallé avec les logiciels dont vous avez besoin dans ce tutoriel, y compris kubectl et gcloud CLI.

  1. Dans la console Google Cloud, démarrez une instance Cloud Shell :
    Ouvrir Cloud Shell

  2. Définissez les variables d'environnement par défaut :

    gcloud config set project PROJECT_ID
    export PROJECT_ID=$(gcloud config get project)
    export LOCATION=LOCATION
    export CLUSTER_NAME=saxml-tpu
    

    Remplacez les valeurs suivantes :

Créer un cluster GKE standard

Dans cette section, vous allez créer le cluster et le pool de nœuds GKE.

Gemma 2B-it

Utilisez Cloud Shell pour effectuer les opérations suivantes :

  1. Créez un cluster Standard qui utilise la fédération d'identité de charge de travail pour GKE :

    gcloud container clusters create ${CLUSTER_NAME} \
        --enable-ip-alias \
        --machine-type=e2-standard-4 \
        --num-nodes=2 \
        --release-channel=rapid \
        --workload-pool=${PROJECT_ID}.svc.id.goog \
        --location=${LOCATION}
    

    La création du cluster peut prendre plusieurs minutes.

  2. Créez un pool de nœuds TPU v5e avec une topologie 1x1 et un nœud:

    gcloud container node-pools create tpu-v5e-1x1 \
        --cluster=${CLUSTER_NAME} \
        --machine-type=ct5lp-hightpu-1t \
        --num-nodes=1 \
        --location=${LOCATION}
    

    Vous diffusez le modèle Gemma 2B dans ce pool de nœuds.

Gemma 7B-it

Utilisez Cloud Shell pour effectuer les opérations suivantes :

  1. Créez un cluster Standard qui utilise la fédération d'identité de charge de travail pour GKE :

    gcloud container clusters create ${CLUSTER_NAME} \
        --enable-ip-alias \
        --machine-type=e2-standard-4 \
        --num-nodes=2 \
        --release-channel=rapid \
        --workload-pool=${PROJECT_ID}.svc.id.goog \
        --location=${LOCATION}
    

    La création du cluster peut prendre plusieurs minutes.

  2. Créez un pool de nœuds TPU v5e avec une topologie 2x2 et un nœud:

    gcloud container node-pools create tpu-v5e-2x2 \
        --cluster=${CLUSTER_NAME} \
        --machine-type=ct5lp-hightpu-4t \
        --num-nodes=1 \
        --location=${LOCATION}
    

    Vous diffusez le modèle Gemma 7B dans ce pool de nœuds.

Créer les buckets Cloud Storage

Créez deux buckets Cloud Storage pour gérer l'état du serveur d'administration Saxml et les points de contrôle du modèle.

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

  1. Créez un bucket Cloud Storage pour stocker les configurations du serveur d'administration Saxml.

    gcloud storage buckets create gs://ADMIN_BUCKET_NAME
    

    Remplacez ADMIN_BUCKET_NAME par le nom du bucket Cloud Storage qui stocke le serveur d'administration Saxml.

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

    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.

Configurer l'accès à vos charges de travail à l'aide de la fédération d'identité de charge de travail pour GKE

Attribuez un compte de service Kubernetes à l'application et configurez ce compte de service Kubernetes pour qu'il agisse en tant que compte de service IAM.

  1. Configurez kubectl de manière à communiquer avec votre cluster :

    gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}
    
  2. Créez un compte de service Kubernetes que votre application pourra utiliser :

    gcloud iam service-accounts create wi-sax
    
  3. Ajoutez une liaison de stratégie IAM pour votre compte de service IAM, pour pouvoir lire et écrire dans Cloud Storage :

    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
        --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
        --role roles/storage.objectUser
    
    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
        --member "serviceAccount:wi-sax@${PROJECT_ID}.iam.gserviceaccount.com" \
        --role roles/storage.insightsCollectorService
    
  4. Autorisez le compte de service Kubernetes à emprunter l'identité du compte de service IAM en ajoutant une liaison de stratégie IAM entre les deux comptes de service. Cette liaison permet au compte de service Kubernetes d'agir en tant que compte de service IAM.

    gcloud iam service-accounts add-iam-policy-binding wi-sax@${PROJECT_ID}.iam.gserviceaccount.com \
        --role roles/iam.workloadIdentityUser \
        --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"
    
  5. Annotez le compte de service Kubernetes avec l'adresse e-mail du compte de service IAM.

    kubectl annotate serviceaccount default \
        iam.gke.io/gcp-service-account=wi-sax@${PROJECT_ID}.iam.gserviceaccount.com
    

Accéder au modèle

Pour accéder aux modèles Gemma en vue du déploiement sur GKE, 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).

Kaggle télécharge un fichier nommé kaggle.json.

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:

  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 un secret Kubernetes pour les identifiants Kaggle

Dans Cloud Shell, procédez comme suit:

  1. Configurez kubectl de manière à communiquer avec votre cluster :

    gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}
    
  2. Créez un secret pour stocker les identifiants Kaggle:

    kubectl create secret generic kaggle-secret \
        --from-file=kaggle.json
    

Déployer Saxml

Dans cette section, vous allez déployer le serveur d'administration Saxml, les serveurs de modèles et le serveur HTTP.

Déployer le serveur d'administration Saxml

  1. Créez le fichier manifeste saxml-admin-server.yaml suivant :

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: sax-admin-server
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: sax-admin-server
      template:
        metadata:
          labels:
            app: sax-admin-server
        spec:
          hostNetwork: false
          containers:
          - name: sax-admin-server
            image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.2.0
            securityContext:
              privileged: true
            ports:
            - containerPort: 10000
            env:
            - name: GSBUCKET
              value: ADMIN_BUCKET_NAME

    Remplacez ADMIN_BUCKET_NAME par le nom du bucket que vous avez créé dans la section Créer des buckets Cloud Storage. N'incluez pas le préfixe gs://.

  2. Appliquez le fichier manifeste :

    kubectl apply -f saxml-admin-server.yaml
    
  3. Vérifiez le déploiement du serveur d'administration:

    kubectl get deployment
    

    La sortie ressemble à ceci :

    NAME                              READY   UP-TO-DATE   AVAILABLE   AGE
    sax-admin-server                  1/1     1            1           ##s
    

Déployer le serveur de modèles Saxml

Suivez ces instructions pour déployer le serveur de modèles pour le modèle Gemma 2B ou Gemma 7B.

Gemma 2B-it

  1. Créez le fichier manifeste saxml-model-server-1x1.yaml suivant :

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: sax-model-server-v5e-1x1
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: gemma-server
      strategy:
        type: Recreate
      template:
        metadata:
          labels:
            app: gemma-server
            ai.gke.io/model: gemma-2b-it
            ai.gke.io/inference-server: saxml
            examples.ai.gke.io/source: user-guide
        spec:
          nodeSelector:
            cloud.google.com/gke-tpu-topology: 1x1
            cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
          hostNetwork: false
          restartPolicy: Always
          containers:
          - name: inference-server
            image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
            args:
            - "--jax_platforms=tpu"
            - "--platform_chip=tpuv5e"
            - "--platform_topology=1x1"
            - "--port=10001"
            - "--sax_cell=/sax/test"
            ports:
            - containerPort: 10001
            securityContext:
              privileged: true
            env:
            - name: SAX_ROOT
              value: "gs://ADMIN_BUCKET_NAME/sax-root"
            resources:
              requests:
                google.com/tpu: 1
              limits:
                google.com/tpu: 1

    Remplacez ADMIN_BUCKET_NAME par le nom du bucket que vous avez créé dans la section Créer des buckets Cloud Storage. N'incluez pas le préfixe gs://.

  2. Appliquez le fichier manifeste :

    kubectl apply -f saxml-model-server-1x1.yaml
    
  3. Vérifiez l'état du déploiement du serveur de modèles:

    kubectl get deployment
    

    La sortie ressemble à ceci :

    NAME                                              READY   STATUS    RESTARTS   AGE
    sax-admin-server                                  1/1     Running   0          ##m
    sax-model-server-v5e-1x1                          1/1     Running   0          ##s
    

Gemma 7B-it

  1. Créez le fichier manifeste saxml-model-server-2x2.yaml suivant :

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: sax-model-server-v5e-2x2
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: gemma-server
      strategy:
        type: Recreate
      template:
        metadata:
          labels:
            app: gemma-server
            ai.gke.io/model: gemma-7b-it
            ai.gke.io/inference-server: saxml
            examples.ai.gke.io/source: user-guide
        spec:
          nodeSelector:
            cloud.google.com/gke-tpu-topology: 2x2
            cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
          hostNetwork: false
          restartPolicy: Always
          containers:
          - name: inference-server
            image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.2.0
            args:
            - "--jax_platforms=tpu"
            - "--platform_chip=tpuv5e"
            - "--platform_topology=2x2"
            - "--port=10001"
            - "--sax_cell=/sax/test"
            ports:
            - containerPort: 10001
            securityContext:
              privileged: true
            env:
            - name: SAX_ROOT
              value: "gs://ADMIN_BUCKET_NAME/sax-root"
            resources:
              requests:
                google.com/tpu: 4
              limits:
                google.com/tpu: 4

    Remplacez ADMIN_BUCKET_NAME par le nom du bucket que vous avez créé dans la section Créer des buckets Cloud Storage. N'incluez pas le préfixe gs://.

  2. Appliquez le fichier manifeste :

    kubectl apply -f saxml-model-server-2x2.yaml
    
  3. Vérifiez l'état du déploiement du serveur de modèles:

    kubectl get deployment
    

    La sortie ressemble à ceci :

    NAME                                              READY   STATUS    RESTARTS   AGE
    sax-admin-server                                  1/1     Running   0          ##m
    sax-model-server-v5e-2x2                          1/1     Running   0          ##s
    

Déployer le serveur HTTP Saxml

Dans cette section, vous allez déployer le serveur HTTP Saxml et créer un service Cluster IP qui vous permettra d'accéder au serveur.

  1. Créez le fichier manifeste saxml-http.yaml suivant :

    apiVersion: apps/v1
    kind: Deployment
    metadata:
      name: sax-http
    spec:
      replicas: 1
      selector:
        matchLabels:
          app: sax-http
      template:
        metadata:
          labels:
            app: sax-http
        spec:
          hostNetwork: false
          containers:
          - name: sax-http
            image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.2.0
            imagePullPolicy: Always
            ports:
            - containerPort: 8888
            env:
            - name: SAX_ROOT
              value: "gs://ADMIN_BUCKET_NAME/sax-root"
    ---
    apiVersion: v1
    kind: Service
    metadata:
      name: sax-http-svc
    spec:
      selector:
        app: sax-http
      ports:
      - protocol: TCP
        port: 8888
        targetPort: 8888
      type: ClusterIP

    Remplacez ADMIN_BUCKET_NAME par le nom du bucket Cloud Storage qui stocke le serveur d'administration Saxml.

  2. Appliquez le fichier manifeste :

    kubectl apply -f saxml-http.yaml
    
  3. Vérifiez l'état du déploiement du serveur HTTP Saxml:

    kubectl get deployment
    

    Gemma 2B-it

    La sortie ressemble à ceci :

      NAME                                              READY   STATUS    RESTARTS   AGE
      sax-admin-server                                  1/1     Running   0          ##m
      sax-model-server-v5e-1x1                          1/1     Running   0          ##m
      sax-http                                          1/1     Running   0          ##s
    

    Gemma 7B-it

    La sortie ressemble à ceci :

      NAME                                              READY   STATUS    RESTARTS   AGE
      sax-admin-server                                  1/1     Running   0          ##m
      sax-model-server-v5e-2x2                          1/1     Running   0          ##m
      sax-http                                          1/1     Running   0          ##s
    

Télécharger le point de contrôle du modèle

Dans cette section, vous allez exécuter une tâche Kubernetes qui récupère, télécharge et stocke le point de contrôle du modèle. Suivez les étapes du modèle Gemma que vous souhaitez utiliser:

Gemma 2B-it

  1. Créez le fichier manifeste job-2b.yaml suivant :

    apiVersion: v1
    kind: ConfigMap
    metadata:
      name: fetch-model-scripts
    data:
      fetch_model.sh: |-
        #!/usr/bin/bash -x
        pip install kaggle --break-system-packages && \
    
        MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
        VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \
    
        mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
        kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
        echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \
    
        gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
        echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: data-loader-2b
      labels:
        app: data-loader-2b
    spec:
      ttlSecondsAfterFinished: 120
      template:
        metadata:
          labels:
            app: data-loader-2b
        spec:
          restartPolicy: OnFailure
          containers:
          - name: gcloud
            image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
            command:
            - /scripts/fetch_model.sh
            env:
            - name: BUCKET_NAME
              value: CHECKPOINTS_BUCKET_NAME
            - name: KAGGLE_CONFIG_DIR
              value: /kaggle
            - name: MODEL_PATH
              value: "google/gemma/pax/2b-it/2"
            volumeMounts:
            - mountPath: "/kaggle/"
              name: kaggle-credentials
              readOnly: true
            - mountPath: "/scripts/"
              name: scripts-volume
              readOnly: true
          volumes:
          - name: kaggle-credentials
            secret:
              defaultMode: 0400
              secretName: kaggle-secret
          - name: scripts-volume
            configMap:
              defaultMode: 0700
              name: fetch-model-scripts

    Remplacez CHECKPOINTS_BUCKET_NAME par le nom du bucket que vous avez créé dans la section Créer des buckets Cloud Storage. N'incluez pas le préfixe gs://.

  2. Appliquez le fichier manifeste :

    kubectl apply -f job-2b.yaml
    
  3. Attendez que la tâche soit terminée :

    kubectl wait --for=condition=complete --timeout=180s job/data-loader-2b
    

    La sortie ressemble à ceci :

    job.batch/data-loader-2b condition met
    
  4. Vérifiez que la tâche a bien été exécutée:

    kubectl get job/data-loader-2b
    

    La sortie ressemble à ceci :

    NAME             COMPLETIONS   DURATION   AGE
    data-loader-2b   1/1           ##s        #m##s
    
  5. Affichez les journaux de la tâche:

    kubectl logs --follow job/data-loader-2b
    

La tâche importe le point de contrôle dans gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000.

Gemma 7B-it

  1. Créez le fichier manifeste job-7b.yaml suivant :

    apiVersion: v1
    kind: ConfigMap
    metadata:
      name: fetch-model-scripts
    data:
      fetch_model.sh: |-
        #!/usr/bin/bash -x
        pip install kaggle --break-system-packages && \
    
        MODEL_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $2}') && \
        VARIATION_NAME=$(echo ${MODEL_PATH} | awk -F'/' '{print $4}') && \
    
        mkdir -p /data/${MODEL_NAME}_${VARIATION_NAME} &&\
        kaggle models instances versions download ${MODEL_PATH} --untar -p /data/${MODEL_NAME}_${VARIATION_NAME} && \
        echo -e "\nCompleted extraction to /data/${MODEL_NAME}_${VARIATION_NAME}" && \
    
        gcloud storage rsync --recursive --no-clobber /data/${MODEL_NAME}_${VARIATION_NAME} gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME} && \
        echo -e "\nCompleted copy of data to gs://${BUCKET_NAME}/${MODEL_NAME}_${VARIATION_NAME}"
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: data-loader-7b
      labels:
        app: data-loader-7b
    spec:
      ttlSecondsAfterFinished: 120
      template:
        metadata:
          labels:
            app: data-loader-7b
        spec:
          restartPolicy: OnFailure
          containers:
          - name: gcloud
            image: gcr.io/google.com/cloudsdktool/google-cloud-cli:slim
            command:
            - /scripts/fetch_model.sh
            env:
            - name: BUCKET_NAME
              value: CHECKPOINTS_BUCKET_NAME
            - name: KAGGLE_CONFIG_DIR
              value: /kaggle
            - name: MODEL_PATH
              value: "google/gemma/pax/7b-it/2"
            volumeMounts:
            - mountPath: "/kaggle/"
              name: kaggle-credentials
              readOnly: true
            - mountPath: "/scripts/"
              name: scripts-volume
              readOnly: true
          volumes:
          - name: kaggle-credentials
            secret:
              defaultMode: 0400
              secretName: kaggle-secret
          - name: scripts-volume
            configMap:
              defaultMode: 0700
              name: fetch-model-scripts

    Remplacez CHECKPOINTS_BUCKET_NAME par le nom du bucket que vous avez créé dans la section Créer des buckets Cloud Storage. Incluez le préfixe gs://.

  2. Appliquez le fichier manifeste :

    kubectl apply -f job-7b.yaml
    
  3. Attendez que la tâche soit terminée :

    kubectl wait --for=condition=complete --timeout=360s job/data-loader-7b
    

    La sortie ressemble à ceci :

    job.batch/data-loader-7b condition met
    
  4. Vérifiez que la tâche a bien été exécutée:

    kubectl get job/data-loader-7b
    

    La sortie ressemble à ceci :

    NAME             COMPLETIONS   DURATION   AGE
    data-loader-7b   1/1           ##s        #m##s
    
  5. Affichez les journaux de la tâche:

    kubectl logs --follow job/data-loader-7b
    

La tâche importe le point de contrôle dans gs://CHECKPOINTS_BUCKET_NAME/gemma_7b_it/checkpoint_00000000.

Exposer le serveur HTTP Saxml

Vous pouvez accéder au serveur HTTP Saxml via le service ClusterIP que vous avez créé lors du déploiement du serveur HTTP Saxml. Les services ClusterIP ne sont accessibles qu'à partir du cluster. Par conséquent, pour accéder au service en dehors du cluster, procédez comme suit:

  1. Établissez une session de transfert de port:

    kubectl port-forward service/sax-http-svc 8888:8888
    
  2. Vérifiez que vous pouvez accéder au serveur HTTP Saxml en ouvrant un nouveau terminal et en exécutant la commande suivante:

    curl -s localhost:8888
    

    La sortie ressemble à ceci :

    {
        "Message": "HTTP Server for SAX Client"
    }
    

Le serveur HTTP Saxml encapsule l'interface client dans le système Saxml et l'expose via un ensemble d'API REST. Vous utilisez ces API pour publier et gérer les modèles Gemma 2B et Gemma 7B et interagir avec eux.

Publier le modèle Gemma

Vous pouvez ensuite publier le modèle Gemma sur un serveur de modèles qui s'exécute dans un pool de nœuds de tranche de pod TPU. Vous utilisez l'API publish du serveur HTTP Saxml pour publier un modèle. Suivez ces étapes pour publier le modèle de paramètre Gemma 2B ou 7B.

Pour en savoir plus sur l'API du serveur HTTP Saxml, consultez la page API HTTP Saxml.

Gemma 2B-it

  1. Assurez-vous que votre session de transfert de port est toujours active:

    curl -s localhost:8888
    
  2. Publiez le paramètre Gemma 2B:

    curl --request POST \
    --header "Content-type: application/json" \
    -s \
    localhost:8888/publish \
    --data \
    '{
        "model": "/sax/test/gemma2bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
        "replicas": "1"
    }'
    

    La sortie ressemble à ceci :

    {
        "model": "/sax/test/gemma2bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
        "replicas": 1
    }
    

    Passez à l'étape suivante pour surveiller la progression du déploiement.

  3. Surveillez la progression en observant les journaux dans un pod de serveur de modèles du déploiement sax-model-server-v5e-1x1.

    kubectl logs --follow deployment/sax-model-server-v5e-1x1
    

    Le déploiement peut prendre jusqu'à cinq minutes. Attendez qu'un message semblable à celui-ci s'affiche:

    I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
    I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma2bfp16
    
  4. Vérifiez que vous pouvez accéder au modèle en affichant les informations associées:

    curl --request GET \
    --header "Content-type: application/json" \
    -s \
    localhost:8888/listcell \
    --data \
    '{
        "model": "/sax/test/gemma2bfp16"
    }'
    

    La sortie ressemble à ceci :

    {
        "model": "/sax/test/gemma2bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma2BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000",
        "max_replicas": 1,
        "active_replicas": 1
    }
    

Gemma 7B-it

  1. Assurez-vous que votre session de transfert de port est toujours active:

    curl -s localhost:8888
    
  2. Publiez le paramètre Gemma 7B:

    curl --request POST \
    --header "Content-type: application/json" \
    -s \
    localhost:8888/publish \
    --data \
    '{
        "model": "/sax/test/gemma7bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
        "replicas": "1"
    }'
    

    La sortie ressemble à ceci :

    {
        "model": "/sax/test/gemma7bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
        "replicas": 1
    }
    

    Passez à l'étape suivante pour surveiller la progression du déploiement.

  3. Surveillez la progression en observant les journaux dans un pod de serveur de modèles du déploiement sax-model-server-v5e-2x2.

    kubectl logs --follow deployment/sax-model-server-v5e-2x2
    

    Attendez qu'un message semblable à celui-ci s'affiche:

    I0125 15:34:31.685555 139063071708736 servable_model.py:699] loading completed.
    I0125 15:34:31.686286 139063071708736 model_service_base.py:532] Successfully loaded model for key: /sax/test/gemma7bfp16
    
  4. Vérifiez que le modèle a été publié en affichant les informations associées:

    curl --request GET \
    --header "Content-type: application/json" \
    -s \
    localhost:8888/listcell \
    --data \
    '{
        "model": "/sax/test/gemma7bfp16"
    }'
    

    Le résultat ressemble à ce qui suit :

    {
        "model": "/sax/test/gemma7bfp16",
        "model_path": "saxml.server.pax.lm.params.gemma.Gemma7BFP16",
        "checkpoint": "gs://CHECKPOINTS_BUCKET_NAME/gemma_7b-it/checkpoint_00000000",
        "max_replicas": 1,
        "active_replicas": 1
    }
    

Utiliser le modèle

Vous pouvez interagir avec les modèles Gemma 2B ou 7B. Utilisez l'API generate du serveur HTTP Saxml pour envoyer une requête au modèle.

Gemma 2B-it

Diffusez une requête en utilisant le point de terminaison generate du serveur HTTP Saxml:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/generate \
--data \
'{
  "model": "/sax/test/gemma2bfp16",
  "query": "What are the top 5 most popular programming languages?"
}'

Voici un exemple de réponse du modèle. Le résultat réel varie en fonction de la requête que vous diffusez:

[
    [
        "\n\n1. **Python**\n2. **JavaScript**\n3. **Java**\n4. **C++**\n5. **Go**",
        -3.0704939365386963
    ]
]

Vous pouvez exécuter la commande avec différents paramètres query. Vous pouvez également modifier des paramètres supplémentaires tels que temperature, top_k et topc_p à l'aide de l'API generate. Pour en savoir plus sur l'API du serveur HTTP Saxml, consultez la page API HTTP Saxml.

Gemma 7B-it

Diffusez une requête en utilisant le point de terminaison generate du serveur HTTP Saxml:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/generate \
--data \
'{
  "model": "/sax/test/gemma7bfp16",
  "query": "What are the top 5 most popular programming languages?"
}'

Voici un exemple de réponse du modèle. Le résultat peut varier pour chaque requête que vous diffusez:

[
    [
        "\n\n**1. JavaScript**\n\n* Most widely used language on the web.\n* Used for front-end development, such as websites and mobile apps.\n* Extensive libraries and frameworks available.\n\n**2. Python**\n\n* Known for its simplicity and readability.\n* Versatile, used for various tasks, including data science, machine learning, and web development.\n* Large and active community.\n\n**3. Java**\n\n* Object-oriented language widely used in enterprise applications.\n* Used for web applications, mobile apps, and enterprise software.\n* Strong ecosystem and support.\n\n**4. Go**\n\n",
        -16.806324005126953
    ]
]

Vous pouvez exécuter la commande avec différents paramètres query. Vous pouvez également modifier des paramètres supplémentaires tels que temperature, top_k et topc_p à l'aide de l'API generate. Pour en savoir plus sur l'API du serveur HTTP Saxml, consultez la page API HTTP Saxml.

Annuler la publication du modèle

Pour annuler la publication de votre modèle, procédez comme suit:

Gemma 2B-it

Pour annuler la publication du modèle Gemma 2B-it, exécutez la commande suivante:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/unpublish \
--data \
'{
    "model": "/sax/test/gemma2bfp16"
}'

La sortie ressemble à ceci :

{
    "model": "/sax/test/gemma2bfp16"
}

Vous pouvez exécuter la commande avec différentes requêtes transmises dans le paramètre query.

Gemma 7B-it

Pour annuler la publication du modèle Gemma 7B-it, exécutez la commande suivante:

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8888/unpublish \
--data \
'{
    "model": "/sax/test/gemma7bfp16"
}'

La sortie ressemble à ceci :

{
    "model": "/sax/test/gemma7bfp16"
}

Vous pouvez exécuter la commande avec différentes requêtes transmises dans le paramètre query.

Résoudre les problèmes

Effectuer un nettoyage

Pour éviter que les ressources utilisées lors de ce tutoriel soient facturées sur votre compte Google Cloud, supprimez le projet contenant les ressources, ou conservez le projet et supprimez les ressources individuelles.

Supprimer les ressources déployées

Pour éviter que les ressources que vous avez créées dans ce guide soient facturées sur votre compte Google Cloud, exécutez la commande suivante:

gcloud container clusters delete ${CLUSTER_NAME} --location=${LOCATION}
gcloud iam service-accounts delete --quiet wi-sax@${PROJECT_ID}.iam.gserviceaccount.com
gcloud storage rm --recursive gs://ADMIN_BUCKET_NAME
gcloud storage rm --recursive gs://CHECKPOINTS_BUCKET_NAME

Remplacez les éléments suivants :

  • ADMIN_BUCKET_NAME: nom du bucket Cloud Storage qui stocke le serveur d'administration Saxml.
  • CHECKPOINTS_BUCKET_NAME: nom du bucket Cloud Storage qui stocke les points de contrôle du modèle.

Étapes suivantes