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

Standard

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.

Expérience

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 :

Découvrez la disponibilité actuelle des versions de TPU avec l'architecture système de Cloud TPU.
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: compose un serveur d'administration et un 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, y compris les points de terminaison pour publier, répertorier, annuler la publication de modèles et générer des prédictions.

Objectifs

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

Ce tutoriel couvre les étapes suivantes :

Préparez un cluster GKE standard avec la topologie TPU recommandée en fonction des caractéristiques du modèle.
Déployer des composants Saxml sur GKE.
Obtenez et publiez le modèle Gemma de paramètre 2B ou 7B.
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 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

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

Activez l'API requise.

Activer l'API

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

Go to project selector

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

Activez l'API requise.

Activer l'API

Assurez-vous que vous disposez du ou des rôles suivants au niveau du projet : roles/container.admin, roles/iam.serviceAccountAdmin
Vérifier les rôles
1. Dans la console Google Cloud, accédez à la page IAM.
  Accéder à IAM
2. Sélectionnez le projet.
3. Dans la colonne Compte principal, recherchez la ligne qui contient votre adresse e-mail.
  
  Si votre adresse e-mail ne figure pas dans cette colonne, cela signifie que vous n'avez aucun rôle.
4. Dans la colonne Rôle de la ligne contenant votre adresse e-mail, vérifiez si la liste des rôles inclut les rôles requis.
Attribuer les rôles
1. Dans la console Google Cloud, accédez à la page IAM.
  Accéder à IAM
2. Sélectionnez le projet.
3. Cliquez sur Accorder l'accès.
4. Dans le champ Nouveaux comptes principaux, saisissez votre adresse e-mail.
5. Dans la liste Sélectinoner 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.

Dans la console Google Cloud, démarrez une instance Cloud Shell :
Ouvrir Cloud Shell
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 :
- PROJECT_ID : L'ID de votre projet Google Cloud.
- LOCATION: nom de la zone Compute Engine où les types de machines TPU v5e sont disponibles.

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 :

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.

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 :

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.

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.

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.
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.

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

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}

Créez un compte de service Kubernetes que votre application pourra utiliser :
```
gcloud iam service-accounts create wi-sax
```

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

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]"
```

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 :

Accédez à la page d'autorisation du modèle sur Kaggle.com.
Connectez-vous à Kaggle si vous ne l'avez 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 un secret Kubernetes pour les identifiants Kaggle

Dans Cloud Shell, procédez comme suit :

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

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${LOCATION}

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

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://.

Appliquez le fichier manifeste :

kubectl apply -f saxml-admin-server.yaml

Vérifiez le déploiement du serveur d'administration:

kubectl get deployment

La sortie devrait ressembler à ce qui suit :

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

Déployer un 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

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://.

Appliquez le fichier manifeste :

kubectl apply -f saxml-model-server-1x1.yaml

Vérifiez l'état du déploiement du serveur de modèles:

kubectl get deployment

La sortie devrait ressembler à ce qui suit :

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

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://.

Appliquez le fichier manifeste :

kubectl apply -f saxml-model-server-2x2.yaml

Vérifiez l'état du déploiement du serveur de modèles:

kubectl get deployment

La sortie devrait ressembler à ce qui suit :

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.

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.

Appliquez le fichier manifeste :
```
kubectl apply -f saxml-http.yaml
```
Vérifiez l'état du déploiement du serveur HTTP Saxml:
```
kubectl get deployment
```

Gemma 2B-it

La sortie devrait ressembler à ce qui suit :

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 devrait ressembler à ce qui suit :

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 en fonction du modèle Gemma que vous souhaitez utiliser:

Gemma 2B-it

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://.

Appliquez le fichier manifeste :
```
kubectl apply -f job-2b.yaml
```

Attendez que la tâche soit terminée :

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

Le résultat ressemble à ce qui suit :

job.batch/data-loader-2b condition met

Vérifiez que la tâche a bien été exécutée:

kubectl get job/data-loader-2b

Le résultat ressemble à ce qui suit :

NAME             COMPLETIONS   DURATION   AGE
data-loader-2b   1/1           ##s        #m##s

Affichez les journaux de la tâche:

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

Les points de contrôle sont importés dans gs://CHECKPOINTS_BUCKET_NAME/gemma_2b-it/checkpoint_00000000.

Gemma 7B-it

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://.

Appliquez le fichier manifeste :
```
kubectl apply -f job-7b.yaml
```

Attendez que la tâche soit terminée :

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

Le résultat ressemble à ce qui suit :

job.batch/data-loader-7b condition met

Vérifiez que la tâche a bien été exécutée:

kubectl get job/data-loader-7b

Le résultat ressemble à ce qui suit :

NAME             COMPLETIONS   DURATION   AGE
data-loader-7b   1/1           ##s        #m##s

Affichez les journaux de la tâche:

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

Les points de contrôle sont importés dans gs://CHECKPOINTS_BUCKET_NAME/gemma_7b_it/checkpoint_00000000.

Exposer le serveur HTTP Saxml

Vous pouvez accéder au serveur HTTP Saxml via un service ClusterIP que vous avez créé à l'étape précédente. 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:

Établissez une session de transfert de port:

kubectl port-forward service/sax-http-svc 8888:8888

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
```
Le résultat ressemble à ce qui suit :
```
{
    "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 utiliserez ces API pour publier et gérer les modèles Gemma 2B et Gemma 7B et interagir avec eux.

Publier le modèle Gemma

Une fois les composants Saxml configurés, vous pouvez 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. Pour en savoir plus sur l'API du serveur HTTP Saxml, consultez la page API HTTP Saxml. Suivez ces étapes pour publier le modèle de paramètre Gemma 2B ou 7B.

Gemma 2B-it

Assurez-vous que votre session de transfert de port est toujours active:
```
curl -s localhost:8888
```

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"
}'

Le résultat ressemble à ce qui suit :

{
    "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
}

Le déploiement du modèle peut prendre quelques minutes.

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

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"
}'

Le résultat ressemble à ce qui suit :

{
    "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

Assurez-vous que votre session de transfert de port est toujours active:
```
curl -s localhost:8888
```

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"
}'

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",
    "replicas": 1
}

Le déploiement du modèle peut prendre quelques minutes.

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

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 invite au modèle.

Gemma 2B-it

Diffusez une requête d'invite 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 peut varier pour chaque invite 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 d'invite 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 invite 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
    ]
]

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"
}'

Le résultat peut ressembler à ce qui suit :

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

Vous pouvez exécuter la commande avec différentes invites 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"
}'

Le résultat peut ressembler à ce qui suit :

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

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

Résoudre les problèmes

Si vous obtenez le message Empty reply from server, il est possible que le conteneur n'ait pas terminé de télécharger les données du modèle. Vérifiez à nouveau dans les journaux du pod. Le message Connected indiquant que le modèle est prêt à être diffusé.
Si Connection refused s'affiche, vérifiez que votre transfert de port est actif.

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

Apprenez-en plus sur les TPU dans GKE.
Explorez le dépôt GitHub Saxml.
Explorez Vertex AI Model Garden.
Découvrez comment exécuter des charges de travail d'IA/de ML optimisées grâce aux fonctionnalités d'orchestration de plates-formes GKE.
Découvrez des architectures de référence, des schémas et des bonnes pratiques concernant Google Cloud. Consultez notre Centre d'architecture cloud.

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

Expérience

Gemma

TPU

Saxml

Objectifs

Architecture

Avant de commencer

Vérifier les rôles

Attribuer les rôles

Préparer l'environnement pour Gemma

Lancer Cloud Shell

Créer un cluster GKE standard

Gemma 2B-it

Gemma 7B-it

Créer les buckets Cloud Storage

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

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

Déployer Saxml

Déployer le serveur d'administration Saxml

Déployer un serveur de modèles Saxml

Gemma 2B-it

Gemma 7B-it

Déployer le serveur HTTP Saxml

Gemma 2B-it

Gemma 7B-it

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

Gemma 2B-it

Gemma 7B-it

Exposer le serveur HTTP Saxml

Publier le modèle Gemma

Gemma 2B-it

Gemma 7B-it

Utiliser le modèle

Gemma 2B-it

Gemma 7B-it

Annuler la publication du modèle

Gemma 2B-it

Gemma 7B-it

Résoudre les problèmes

Effectuer un nettoyage

Supprimer les ressources déployées

Étapes suivantes