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.

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 :

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

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, tels que ceux créés dans ce tutoriel.

Avant de commencer

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

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

Go to project selector

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

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

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

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:

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

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

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

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

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

La sortie ressemble à ceci :

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

La sortie ressemble à ceci :

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

La tâche importe le point de contrôle 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

La sortie ressemble à ceci :

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

La sortie ressemble à ceci :

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

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:

É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
```
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

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

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.

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

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

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

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.

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

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

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, y compris les API HTTP 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 Cloud Architecture Center.

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

Contexte

Gemma

TPU

Saxml

Objectifs

Architecture

Avant de commencer

Check for the roles

Grant the roles

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