Diffuser un LLM à l'aide de TPU multi-hôtes sur GKE avec Saxml

Ce tutoriel explique comment diffuser un grand modèle de langage (LLM) à l'aide des TPU (Tensor Processing Units) sur Google Kubernetes Engine (GKE) avec Saxml.

Expérience

Saxml est un système expérimental qui diffuse les frameworks Paxml, JAX et PyTorch. Vous pouvez utiliser les TPU pour accélérer le traitement des données avec ces frameworks. Pour illustrer le déploiement de TPU dans GKE, ce tutoriel diffuse le modèle de test LmCloudSpmd175B32Test 175B. GKE déploie ce modèle de test sur deux pools de nœuds TPU v5e avec la topologie 4x8 respectivement.

Pour déployer correctement le modèle de test, la topologie TPU a été définie en fonction de la taille du modèle. Étant donné que le modèle de N milliard 16 bits nécessite environ deux fois (2xN) Go de mémoire, le modèle LmCloudSpmd175B32Test 175B nécessite environ 350 Go de mémoire. La puce TPU v5e dispose de 16 Go de mémoire. Pour accepter 350 Go, GKE a besoin de 21 puces v5e (350/16= 21). En fonction du mappage de la configuration TPU, la configuration TPU appropriée pour ce tutoriel est la suivante:

  • Type de machine : ct5lp-hightpu-4t
  • Topologie: 4x8 (32 nombre de puces TPU)

Il est important de sélectionner la topologie TPU appropriée pour diffuser un modèle lors du déploiement de TPU dans GKE. Pour en savoir plus, consultez la section Planifier la configuration de TPU.

Objectifs

Ce tutoriel est destiné aux ingénieurs MLOps ou DevOps ou aux administrateurs de plate-forme qui souhaitent utiliser les fonctionnalités d'orchestration GKE pour diffuser des modèles de données.

Ce tutoriel couvre les étapes suivantes :

  1. Préparez votre environnement avec un cluster GKE standard. Le cluster comporte deux pools de nœuds TPU v5e avec la topologie 4x8.
  2. Déployer Saxml. Saxml nécessite un serveur d'administration, un groupe de pods qui fonctionnent en tant que serveur de modèles, un serveur HTTP prédéfini et un équilibreur de charge.
  3. Utilisez le fichier Saxml pour diffuser le LLM.

Le schéma suivant illustre l'architecture mise en œuvre dans le tutoriel suivant:

Architecture d'un TPU multi-hôte sur GKE.
Figure: Exemple d'architecture d'un TPU multi-hôte sur GKE.

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.

Préparer l'environnement

  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 REGION=COMPUTE_REGION
  export ZONE=COMPUTE_ZONE
  export GSBUCKET=PROJECT_ID-gke-bucket

    Remplacez les valeurs suivantes :

Créer un cluster GKE standard

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 saxml \
    --zone=${ZONE} \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --cluster-version=VERSION \
    --num-nodes=4

    Remplacez VERSION par le numéro de version de GKE. GKE est compatible avec les TPU v5e version 1.27.2-gke.2100 et ultérieures. Pour en savoir plus, consultez la section Disponibilité des TPU dans GKE.

    La création du cluster peut prendre plusieurs minutes.

  2. Créez le premier pool de nœuds nommé tpu1:

    gcloud container node-pools create tpu1 \
    --zone=${ZONE} \
    --num-nodes=8 \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=4x8 \
    --cluster=saxml

  3. Créez le deuxième pool de nœuds nommé tpu2:

    gcloud container node-pools create tpu2 \
    --zone=${ZONE} \
    --num-nodes=8 \
    --machine-type=ct5lp-hightpu-4t \
    --tpu-topology=4x8 \
    --cluster=saxml

Vous avez créé les ressources suivantes:

  • Cluster standard avec quatre nœuds de processeur.
  • Deux pools de nœuds TPU v5e avec la topologie 4x8. Chaque pool de nœuds représente huit nœuds TPU avec quatre puces chacun.

Le modèle 175B doit être diffusé sur une tranche de TPU v5e à plusieurs hôtes avec une tranche de topologie 4x8 (32 puces TPU v5e) au minimum.

Créer un bucket Cloud Storage

Créez un bucket Cloud Storage pour stocker les configurations du serveur d'administration Saxml. Un serveur d'administration en cours d'exécution enregistre régulièrement son état et les détails des modèles publiés.

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

gcloud storage buckets create gs://${GSBUCKET}

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 saxml --zone=${ZONE}

  2. Créez un compte de service Kubernetes que votre application pourra utiliser :

    kubectl create serviceaccount sax-sa --namespace default

  3. Créez un compte de service IAM pour votre application :

    gcloud iam service-accounts create sax-iam-sa

  4. 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:sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com" \
  --role roles/storage.admin

  5. 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, afin que le compte de service Kubernetes puisse lire et écrire dans Cloud Storage.

    gcloud iam service-accounts add-iam-policy-binding sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:${PROJECT_ID}.svc.id.goog[default/sax-sa]"

  6. Annotez le compte de service Kubernetes avec l'adresse e-mail du compte de service IAM. Cela permet à votre exemple d'application d'identifier le compte de service à utiliser pour accéder aux services Google Cloud. Ainsi, lorsque l'application utilise des bibliothèques clientes d'API Google standards pour accéder aux services Google Cloud, elle utilise ce compte de service IAM.

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

Déployer Saxml

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

Déployer le serveur d'administration Saxml

  1. Créez le fichier manifeste sax-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
      serviceAccountName: sax-sa
      containers:
      - name: sax-admin-server
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.1.0
        securityContext:
          privileged: true
        ports:
        - containerPort: 10000
        env:
        - name: GSBUCKET
          value: BUCKET_NAME

    Remplacez BUCKET_NAME par le nom de votre bucket Cloud Storage.

  2. Appliquez le fichier manifeste :

    kubectl apply -f sax-admin-server.yaml

  3. Vérifiez que le pod du serveur d'administration est opérationnel:

    kubectl get deployment

    Le résultat ressemble à ce qui suit :

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

Déployer un serveur de modèles Saxml

Les charges de travail exécutées dans des tranches de TPU multi-hôte nécessitent un identifiant réseau stable pour que chaque pod puisse découvrir les pairs dans la même tranche de TPU. Pour définir ces identifiants, utilisez IndexedJob, StatefulSet avec un service sans adresse IP ou JobSet qui crée automatiquement un service sans adresse IP pour tous les jobs appartenant à JobSet. La section suivante explique comment gérer plusieurs groupes de pods de serveur de modèles avec JobSet.

  1. Installez JobSet v0.2.3 ou version ultérieure.

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml

    Remplacez JOBSET_VERSION par la version de JobSet. Par exemple, v0.2.3.

  2. Vérifiez que le contrôleur JobSet s'exécute dans l'espace de noms jobset-system:

    kubectl get pod -n jobset-system

    Le résultat ressemble à ce qui suit :

    NAME                                        READY   STATUS    RESTARTS   AGE
jobset-controller-manager-69449d86bc-hp5r6   2/2     Running   0          2m15s

  3. Déployer deux serveurs de modèles dans deux pools de nœuds TPU. Enregistrez le fichier manifeste sax-model-server-set suivant :

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: sax-model-server-set
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: sax-model-server
      replicas: 2
      template:
        spec:
          parallelism: 8
          completions: 8
          backoffLimit: 0
          template:
            spec:
              serviceAccountName: sax-sa
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 4x8
              containers:
              - name: sax-model-server
                image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0
                args: ["--port=10001","--sax_cell=/sax/test", "--platform_chip=tpuv5e"]
                ports:
                - containerPort: 10001
                - containerPort: 8471
                securityContext:
                  privileged: true
                env:
                - name: SAX_ROOT
                  value: "gs://BUCKET_NAME/sax-root"
                - name: MEGASCALE_NUM_SLICES
                  value: ""
                resources:
                  requests:
                    google.com/tpu: 4
                  limits:
                    google.com/tpu: 4

    Remplacez BUCKET_NAME par le nom de votre bucket Cloud Storage.

    Dans le fichier manifeste :

    • replicas: 2 est le nombre d'instances dupliquées de tâches. Chaque tâche représente un serveur de modèles. Par conséquent, un groupe de 8 pods.
    • parallelism: 8 et completions: 8 sont égaux au nombre de nœuds de chaque pool de nœuds.
    • backoffLimit: 0 doit être égal à zéro pour marquer la tâche comme ayant échoué en cas de défaillance d'un pod.
    • ports.containerPort: 8471 est le port par défaut pour la communication des VM TPU.
    • name: MEGASCALE_NUM_SLICES annule la définition de la variable d'environnement, car GKE n'exécute pas d'entraînement Multislice.

  4. Appliquez le fichier manifeste :

    kubectl apply -f sax-model-server-set.yaml

  5. Vérifiez l'état des pods du serveur d'administration Saxml et du serveur de modèles:

    kubectl get pods

    Le résultat ressemble à ce qui suit :

    NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server-557c85f488-lnd5d                 1/1     Running   0          35h
sax-model-server-set-sax-model-server-0-0-nj4sm   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-1-sl8w4   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-2-hb4rk   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-3-qv67g   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-4-pzqz6   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-5-nm7mz   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-6-7br2x   1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-7-4pw6z   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-0-8mlf5   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-1-h6z6w   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-2-jggtv   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-3-9v8kj   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-4-6vlb2   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-5-h689p   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-6-bgv5k   1/1     Running   0          24m
sax-model-server-set-sax-model-server-1-7-cd6gv   1/1     Running   0          24m

Dans cet exemple, il existe 16 conteneurs de serveurs de modèles : sax-model-server-set-sax-model-server-0-0-nj4sm et sax-model-server-set-sax-model-server-1-0-8mlf5 sont les deux serveurs de modèles principaux dans chaque groupe.

Votre cluster Saxml possède deux serveurs de modèles déployés sur deux pools de nœuds TPU v5e avec la topologie 4x8 respectivement.

Déployer Saxml HTTP Server et l'équilibreur de charge

  1. Utilisez l'image du serveur HTTP d'images prédéfinies suivante. Enregistrez le fichier manifeste sax-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
      serviceAccountName: sax-sa
      containers:
      - name: sax-http
        image: us-docker.pkg.dev/cloud-tpu-images/inference/sax-http:v1.0.0
        ports:
        - containerPort: 8888
        env:
        - name: SAX_ROOT
          value: "gs://BUCKET_NAME/sax-root"
---
apiVersion: v1
kind: Service
metadata:
  name: sax-http-lb
spec:
  selector:
    app: sax-http
  ports:
  - protocol: TCP
    port: 8888
    targetPort: 8888
  type: LoadBalancer

    Remplacez BUCKET_NAME par le nom de votre bucket Cloud Storage.

  2. Appliquez le fichier manifeste sax-http.yaml :

    kubectl apply -f sax-http.yaml

  3. Attendez la fin de la création du conteneur de serveur HTTP:

    kubectl get pods

    Le résultat ressemble à ce qui suit :

    NAME                                              READY   STATUS    RESTARTS   AGE
sax-admin-server-557c85f488-lnd5d                 1/1     Running   0          35h
sax-http-65d478d987-6q7zd                         1/1     Running   0          24m
sax-model-server-set-sax-model-server-0-0-nj4sm   1/1     Running   0          24m
...

  4. Attendez qu'une adresse IP externe soit attribuée au service:

    kubectl get svc

    Le résultat ressemble à ce qui suit :

    NAME           TYPE           CLUSTER-IP    EXTERNAL-IP   PORT(S)          AGE
sax-http-lb    LoadBalancer   10.48.11.80   10.182.0.87   8888:32674/TCP   7m36s

Utiliser Saxml

Chargez, déployez et diffusez le modèle sur Saxml dans la tranche multi-hôte TPU v5e:

Charger le modèle

  1. Récupérez l'adresse IP de l'équilibreur de charge pour Saxml.

    LB_IP=$(kubectl get svc sax-http-lb -o jsonpath='{.status.loadBalancer.ingress[*].ip}')
PORT="8888"

  2. Chargez le modèle de test LmCloudSpmd175B dans deux pools de nœuds TPU v5e:

    curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/publish --data \
'{
    "model": "/sax/test/spmd",
    "model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
    "checkpoint": "None",
    "replicas": 2
}'

    Le modèle de test n'a pas de point de contrôle affiné. Les pondérations sont générées de manière aléatoire. Le chargement du modèle peut prendre jusqu'à 10 minutes.

    Le résultat ressemble à ce qui suit :

    {
    "model": "/sax/test/spmd",
    "path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
    "checkpoint": "None",
    "replicas": 2
}

  3. Vérifiez que le modèle est prêt:

    kubectl logs sax-model-server-set-sax-model-server-0-0-nj4sm

    Le résultat ressemble à ce qui suit :

    ...
loading completed.
Successfully loaded model for key: /sax/test/spmd

    Le modèle est entièrement chargé.

  4. Obtenez des informations sur le modèle:

    curl --request GET \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/listcell --data \
'{
    "model": "/sax/test/spmd"
}'

    Le résultat ressemble à ce qui suit :

    {
"model": "/sax/test/spmd",
"model_path": "saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test",
"checkpoint": "None",
"max_replicas": 2,
"active_replicas": 2
}

Diffuser le modèle

Diffuser une requête d'invite:

curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/generate --data \
'{
  "model": "/sax/test/spmd",
  "query": "How many days are in a week?"
}'

Le résultat affiche un exemple de réponse du modèle. Cette réponse peut ne pas avoir de sens, car le modèle de test a des pondérations aléatoires.

Annuler la publication du modèle

Exécutez la commande suivante pour annuler la publication du modèle:

curl --request POST \
--header "Content-type: application/json" \
-s ${LB_IP}:${PORT}/unpublish --data \
'{
    "model": "/sax/test/spmd"
}'

Le résultat ressemble à ce qui suit :

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

Nettoyer

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

  1. Supprimez le cluster que vous avez créé pour ce tutoriel:

    gcloud container clusters delete saxml --zone ${ZONE}

  2. Supprimez le compte de service :

    gcloud iam service-accounts delete sax-iam-sa@${PROJECT_ID}.iam.gserviceaccount.com

  3. Supprimez le bucket Cloud Storage :

    gcloud storage rm -r gs://${GSBUCKET}

Étapes suivantes