Diffuser des LLM à l'aide de TPU multi-hôtes sur GKE avec JetStream et Pathways

Standard

---

Ce guide vous explique comment diffuser des grands modèles de langage (LLM) de pointe tels que Llama 3.1 405B sur Google Kubernetes Engine (GKE) à l'aide d'unités de traitement tensoriel (TPU) sur plusieurs nœuds.

Ce guide explique comment utiliser des technologies Open Source portables (Kubernetes, JetStream, Pathways on Cloud et l'API LeaderWorkerSet (LWS)) pour déployer et diffuser des charges de travail d'IA/de ML sur GKE, en tirant parti du contrôle précis, de l'évolutivité, de la résilience, de la portabilité et de la rentabilité de GKE.

Important : Pour accéder aux parcours, contactez votre responsable de compte Google Cloud.

Arrière-plan

Les grands modèles de langage ont augmenté en taille et ne tiennent plus sur une seule tranche de TPU hôte. Pour l'inférence ML, vous pouvez utiliser Pathways on Cloud afin d'exécuter une inférence multihôte à grande échelle sur GKE sur plusieurs nœuds TPU interconnectés. Dans ce guide, vous allez apprendre à provisionner un cluster GKE avec les tranches TPU multi-hôtes, à utiliser les binaires Pathways on Cloud, à lancer le serveur JetStream avec le framework MaxText et à effectuer des requêtes d'inférence multi-hôtes.

En diffusant un LLM à l'aide de TPU sur GKE avec JetStream, MaxText et Pathways, vous pouvez créer une solution de diffusion robuste et prête pour la production avec tous les avantages de la plate-forme Kubernetes gérée, y compris en termes de rentabilité, évolutivité et haute disponibilité. Cette section décrit les principales technologies utilisées dans ce tutoriel.

À propos des 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 le machine learning et les modèles d'IA créés à l'aide de frameworks tels que TensorFlow, PyTorch et JAX.

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

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

Ce tutoriel explique comment diffuser le modèle Llama 3.1-405B. GKE déploie le modèle sur des nœuds TPU v6e multi-hôte avec des topologies TPU configurées en fonction des exigences du modèle pour diffuser des requêtes avec une faible latence.

Parcours sur Cloud

Pathways est une couche d'orchestration à grande échelle pour les accélérateurs. Pathways est explicitement conçu pour permettre l'exploration de nouveaux systèmes et de nouvelles idées de recherche en ML, tout en conservant des performances de pointe pour les modèles actuels. Pathways permet à un seul processus client JAX de coordonner le calcul sur une ou plusieurs grandes tranches de TPU, ce qui simplifie les calculs de ML qui s'étendent sur des centaines ou des milliers de puces TPU.

JetStream

JetStream est un framework de diffusion d'inférences Open Source développé par Google. JetStream permet des inférences hautes performances, à haut débit et à mémoire optimisée sur les TPU et les GPU. JetStream fournit des optimisations de performances avancées, y compris des techniques de traitement par lot continu, d'optimisation du cache KV et de quantification, pour faciliter le déploiement de LLM. JetStream permet d'optimiser les performances de diffusion des TPU PyTorch/XLA et JAX.

MaxText

MaxText est une implémentation LLM JAX performante, évolutive et adaptable, basée sur des bibliothèques JAX Open Source telles que Flax, Orbax et Optax. L'implémentation LLM uniquement décodeur de MaxText est écrite en Python. Elle exploite fortement le compilateur XLA pour atteindre de hautes performances sans avoir à créer de noyau personnalisé.

Pour en savoir plus sur les derniers modèles et tailles de paramètres compatibles avec MaxText, consultez le dépôt du projet MaxText.

Llama 3.1 405B

Llama 3.1 405B est un grand modèle de langage de Meta conçu pour diverses tâches de traitement du langage naturel, y compris la génération de texte, la traduction et les systèmes de questions-réponses. GKE fournit l'infrastructure nécessaire pour répondre aux besoins d'entraînement et de diffusion distribués des modèles de cette envergure.

Pour en savoir plus, consultez la documentation Llama.

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 JetStream et Pathways pour déployer et diffuser le modèle.

Le schéma suivant montre les composants de cette architecture:

Cette architecture comprend les composants suivants :

• Un cluster régional GKE Standard.
• Un pool de nœuds de tranche de TPU multi-hôte qui héberge le déploiement JetStream et les composants Pathways.
• Pathways resource manager gère les ressources d'accélérateur et coordonne l'allocation des accélérateurs pour les jobs utilisateur.
• Pathways client se coordonne avec Pathways resource manager pour déterminer où les programmes compilés sont placés pour l'exécution.
• Le Pathways worker s'exécute et effectue des calculs sur les machines d'accélérateur, puis renvoie les données à votre charge de travail via le serveur proxy IFRT.
• IFRT proxy client implémente l'API Interim Framework Runtime (IFRT) OSS et sert de pont de communication entre votre charge de travail et les composants Pathways.
• Le IFRT proxy server reçoit les requêtes du IFRT proxy client et les transmet au Pathways client, en répartissant le travail.
• Le conteneur JetStream-Pathways fournit un serveur d'inférence basé sur JAX qui reçoit les requêtes d'inférence et délègue ses processus d'exécution au Pathways workers.
• Le composant de service répartit le trafic entrant sur toutes les instances répliquées JetStream HTTP.
• JetStream HTTP est un serveur HTTP qui accepte les requêtes en tant qu'encapsuleur au format requis de JetStream et les envoie au client GRPC de JetStream.

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.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the 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.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the 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, roles/resourcemanager.projectIamAdmin

  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 column 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 d'un quota suffisant pour seize puces TPU v6e PodSlice Lite. Dans ce tutoriel, vous utilisez des instances à la demande.
• Assurez-vous que votre projet Google Cloud figure dans la liste d'autorisation pour Pathways.

Accéder au modèle

Pour accéder au point de contrôle Meta Llama 3.1-405B afin de le déployer sur GKE, procédez comme suit :

1. Signez le contrat de consentement de la licence.
2. Accédez à la page de téléchargement de Meta Llama.
3. Lisez et acceptez les Conditions d'utilisation du modèle pour obtenir l'URL nécessaire à son téléchargement.
4. Pour télécharger le point de contrôle du modèle, recherchez l'ID du modèle approprié. Pour obtenir la liste des modèles compatibles et de leurs ID, consultez la documentation de la CLI Llama. Par exemple, utilisez Llama 3.1-405B-Instruct:bf16-mp16 pour le modèle Llama 3.1-405B.

Préparer l'environnement

Dans ce tutoriel, vous utilisez Cloud Shell pour gérer les ressources hébergées surGoogle Cloud. Les logiciels dont vous avez besoin pour ce tutoriel sont préinstallés sur Cloud Shell, y compris kubectl et gcloud CLI.

Pour configurer votre environnement avec Cloud Shell, procédez comme suit :

1. Dans la console Google Cloud , lancez une session Cloud Shell en cliquant sur Activer Cloud Shell dans la consoleGoogle Cloud . Une session s'ouvre dans le volet inférieur de la console Google Cloud .

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

   gcloud config set project PROJECT_ID
   gcloud config set billing/quota_project PROJECT_ID
   export PROJECT_ID=$(gcloud config get project)
   export CLUSTER_NAME=CLUSTER_NAME
   export BUCKET_NAME=BUCKET_NAME
   export REGION=REGION
   export LOCATION=LOCATION
   export CLUSTER_VERSION=CLUSTER_VERSION
   export MACHINE_TYPE=ct6e-standard-4t
   export TPU_TYPE=v6e
   export TOPOLOGY=4x4
   export WORKERS_PER_SLICE=4
   

   Remplacez les valeurs suivantes :

   • PROJECT_ID : ID de votre projet Google Cloud .
   • CLUSTER_NAME : nom de votre cluster GKE.
   • BUCKET_NAME : nom du bucket Cloud Storage. Vous n'avez pas besoin de spécifier le préfixe gs://.
   • REGION: région où se trouvent le cluster GKE, le bucket Cloud Storage et les nœuds TPU. La région contient des zones où les types de machines TPU v6e sont disponibles (par exemple, us-east1, us-east5, europe-west4, asia-northeast1 ou us-south1).
   • LOCATION : zone dans laquelle les ressources TPU sont disponibles (par exemple, us-east1-d).
   • CLUSTER_VERSION : version de GKE, qui doit être compatible avec le type de machine que vous souhaitez utiliser. Notez que la version de GKE par défaut peut ne pas avoir de disponibilité pour votre TPU cible. Pour obtenir la liste des versions minimales de GKE disponibles par type de machine TPU, consultez Disponibilité des TPU dans GKE.
   • MACHINE_TYPE : type de machine v6e.
   • TPU_TYPE : préfixe utilisé pour nommer les pools de nœuds (v6e).
   • TOPOLOGY : topologie TPU v6e.
   • WORKERS_PER_SLICE : nombre de nœuds par pool de nœuds ou tranche TPU.

Créer et configurer des ressources Google Cloud

Pour créer les ressources requises, suivez ces instructions :

Remarque : Vous devrez peut-être créer une réservation de capacité pour utiliser certains accélérateurs. Pour savoir comment réserver et consommer des ressources réservées, consultez la section Consommer des ressources zonales réservées.

Créer un cluster GKE

1. Créez un cluster GKE Standard régional :

   gcloud container clusters create CLUSTER_NAME \
       --project=PROJECT_ID \
       --cluster-version=CLUSTER_VERSION \
       --region=REGION \
       --scopes=cloud-platform \
       --machine-type=n2-standard-32
   

   La création du cluster peut prendre plusieurs minutes.

   Remplacez CLUSTER_VERSION par la version du cluster appropriée.

2. Créez un pool de nœuds TPU v6e avec une topologie 4x4 et quatre nœuds chacun :

   gcloud container node-pools create multihost-np \
   --project=PROJECT_ID \
   --region=REGION \
   --node-locations=LOCATION \
   --cluster=CLUSTER_NAME \
   --machine-type=MACHINE_TYPE \
   --num-nodes=WORKERS_PER_SLICE \
   --tpu-topology=TOPOLOGY \
   --scopes cloud-platform \
   --placement-type=COMPACT \
   --workload-metadata=GCE_METADATA
   

   Remarque : Le symbole --workload-metadata=GCE_METADATA indique que la fédération d'identité de charge de travail pour GKE est désactivée pour le pool de nœuds. Les pools de nœuds qui exécutent Pathways avec JetStream et MaxText doivent désactiver la fédération d'identité de charge de travail pour GKE.

Configurer un compte de service pour l'accès aux objets Storage

Configurez un compte de service Kubernetes qui agira en tant que compte de service IAM.

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

   gcloud iam service-accounts create jetstream-pathways
   

2. Ajoutez une liaison de stratégie IAM pour votre compte de service IAM afin de gérer Cloud Storage. Cela permet à votre compte de service IAM d'accéder au bucket de stockage dans lequel votre point de contrôle sera stocké :

   gcloud projects add-iam-policy-binding ${PROJECT} \
     --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
     --role roles/storage.objectUser
   
   gcloud projects add-iam-policy-binding ${PROJECT} \
     --member "serviceAccount:jetstream-pathways@${PROJECT}.iam.gserviceaccount.com" \
     --role roles/storage.insightsCollectorService
   

3. 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=jetstream-pathways@${PROJECT}.iam.gserviceaccount.com
   

Configurer Docker pour l'authentification auprès d'Artifact Registry

Configurez Docker pour qu'il s'authentifie auprès d'Artifact Registry afin de pouvoir extraire les images Pathways figurant sur la liste d'autorisation :

gcloud auth login
gcloud auth configure-docker

Conversion de point de contrôle

Pour convertir un point de contrôle Meta Llama 3.1-405B en point de contrôle d'inférence int8 compatible avec MaxText, suivez les étapes décrites dans Conversion de points de contrôle avec Llama3.1-405B. Votre déploiement utilise le point de contrôle avec l'indicateur load_parameters_path.

Créer un bucket Cloud Storage pour stocker les fichiers temporaires Pathways

Créez un bucket Cloud Storage pour stocker vos fichiers temporaires Pathways, comme le cache de compilation :

export PATHWAYS_BUCKET=PATHWAYS_BUCKET
gcloud storage buckets create gs://$PATHWAYS_BUCKET

Déployer JetStream-MaxText et Pathways

Déployez le serveur de modèles JetStream-MaxText et Pathways.

Se connecter au cluster GKE

gcloud container clusters get-credentials "${CLUSTER}" --project "${PROJECT}" --location "${ZONE}"

Déployer l'API LeaderWorkerSet (LWS)

LWS est une ressource personnalisée conçue pour déployer et gérer des applications distribuées avec état, en particulier celles dotées d'une architecture leader-worker. Il est particulièrement adapté aux charges de travail d'IA/ML où un grand modèle est fragmenté et diffusé sur plusieurs appareils et nœuds.

VERSION=v0.6.1
kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/$VERSION/manifests.yaml

Attendez que le contrôleur LeaderWorkerSet soit entièrement disponible :

kubectl wait deploy/lws-controller-manager -n lws-system --for=condition=available --timeout=5m

La sortie devrait ressembler à ce qui suit :

deployment.apps/lws-controller-manager condition met

Vérifiez que le contrôleur LeaderWorkerSet s'exécute dans l'espace de noms lws-system :

kubectl get pod -n lws-system

La sortie devrait ressembler à ce qui suit :

NAME                          READY   STATUS    RESTARTS    AGE
lws-controller-manager-abcd   1/1     Running   0           40s
lws-controller-manager-efgh   1/1     Running   0           40s

Déployer le fichier manifeste de la charge de travail

1. Enregistrez le manifeste suivant sous le nom jetstream-pathways-llama-3-1-405b-4x4.yaml :

   apiVersion: leaderworkerset.x-k8s.io/v1
   kind: LeaderWorkerSet
   metadata:
     name: jetstream-pathways
     annotations:
       leaderworkerset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     replicas: 1
     leaderWorkerTemplate:
       leaderTemplate:
         metadata:
           labels:
             app: jetstream-pathways
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 4x4
           tolerations:
           - key: "google.com/tpu"
             operator: "Exists"
             effect: "NoSchedule"
           containers:
           - name: pathways-proxy
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
             args:
             imagePullPolicy: Always
             ports:
             - containerPort: 38681
           - name: pathways-rm
             env:
             - name: HOST_ADDRESS
               value: "$(LWS_LEADER_ADDRESS)"
             - name: TPU_SKIP_MDS_QUERY
               value: "true"
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             args:
             - --server_port=38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             - --node_type=resource_manager
             - --instance_count=1
             - --instance_type=tpuv6e:4x4
             imagePullPolicy: Always
             ports:
             - containerPort: 38677
           - name: jax-tpu
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
             env:
             - name: LOG_LEVEL
               value: "INFO"
             args:
             - MaxText/configs/v5e/inference/llama3_405b_v5e-64.yml
             - model_name=llama3.1-405b
             - load_parameters_path=CHECKPOINT_PATH
             - max_prefill_predict_length=1024
             - max_target_length=2048
             - async_checkpointing=false
             - steps=1
             - ici_fsdp_parallelism=1
             - ici_autoregressive_parallelism=2
             - ici_tensor_parallelism=8
             - scan_layers=false
             - weight_dtype=bfloat16
             - per_device_batch_size=6
             - enable_single_controller=true
             - quantization=int8
             - quantize_kvcache=true
             - checkpoint_is_quantized=true
             - enable_model_warmup=true
             imagePullPolicy: Always
             ports:
             - containerPort: 9000
             startupProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 1
               initialDelaySeconds: 600
               failureThreshold: 10000
             livenessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 10
             readinessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 10
           - name: jetstream-http
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
             imagePullPolicy: Always
             ports:
             - containerPort: 8000
       size: 5
       workerTemplate:
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 4x4
           tolerations:
           - key: "google.com/tpu"
             operator: "Exists"
             effect: "NoSchedule"
           containers:
           - name: worker
             args:
             - --server_port=38679
             - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             imagePullPolicy: Always
             ports:
             - containerPort: 38679
             resources:
               limits:
                 google.com/tpu: "4"
   --- 
   apiVersion: v1
   kind: Service
   metadata:
     name: jetstream-svc
   spec:
     selector:
       app: jetstream-pathways
     ports:
     - protocol: TCP
       name: jetstream-http
       port: 8000
       targetPort: 8000

2. Définissez la valeur du champ load_parameters_path sur le chemin d'accès de votre point de contrôle qui a été généré lors du processus de conversion des points de contrôle.

   • Pour un point de contrôle bf16, le chemin d'accès doit être semblable à gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
   • Pour un point de contrôle int8, il doit être semblable à gs://OUTPUT_BUCKET_DIRECTORY/int8.

   Définissez la valeur du champ gcs_scratch_location sur le bucket Pathways que vous avez créé précédemment.

   perl -pi -e 's|CHECKPOINT_PATH|gs://OUTPUT_BUCKET_DIRECTORY/int8|g' jetstream-pathways-llama-3-1-405b-4x4.yaml
   perl -pi -e 's|PATHWAYS_BUCKET|gs://PATHWAYS_BUCKET|g' jetstream-pathways-llama-3-1-405b-4x4.yaml
   

   Description des variables et des paramètres.

   Les descriptions suivantes expliquent les variables et les paramètres pertinents pour JetStream, Pathways et LeaderWorkerSet.

   • Jetstream
     • Ce conteneur est créé à partir du projet OSS JetStream. Il encapsule également la bibliothèque cliente basée sur gRPC du proxy IFRT, qui fournit un backend JAX pouvant déléguer l'exécution de programmes TPU au serveur proxy Pathways.
     • Pour en savoir plus sur les indicateurs JetStream-MaxText, consultez Descriptions des indicateurs du serveur JetStream-MaxText.
     • Pour utiliser un point de contrôle quantifié, vous pouvez définir les options quantization, quantize_kvcache ou checkpoint_is_quantized. Si le point de contrôle est bf16, vous pouvez omettre ces indicateurs.
   • Parcours : pour obtenir une description des composants de parcours concernés, consultez Composants de parcours.
   • LeaderWorkerSet : pour en savoir plus sur les paramètres LeaderWorkerSet, consultez la définition de l'API.

Appliquer le fichier manifeste de déploiement

Appliquez le fichier manifeste pour déployer le serveur :

kubectl apply -f jetstream-pathways-llama-3-1-405b-4x4.yaml

Le serveur de modèles devrait démarrer.

Vérifier le démarrage du serveur de modèles

La restauration du point de contrôle peut prendre entre 10 et 20 minutes pour un modèle 405B. Vous pouvez également attendre plus longtemps pendant l'échauffement du modèle si vous avez activé l'indicateur enable_model_warmup.

kubectl logs -f jetstream-pathways-0 -c jax-tpu

Le résultat ressemble à ce qui suit :

2025-03-02 02:15:07,682 - JetstreamLogger - INFO - Initializing the driver with 1 prefill engines and 1 generate engines in interleaved mode
2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up prefill thread 0.
2025-03-02 02:15:07,683 - JetstreamLogger - INFO - Spinning up transfer thread 0.
2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up generate thread 0.
2025-03-02 02:15:07,684 - JetstreamLogger - INFO - Spinning up detokenize thread 0.
2025-03-02 02:15:07,685 - JetstreamLogger - INFO - Driver initialized.
...
...
...
INFO:     Started server process [7]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:9999 (Press CTRL+C to quit)

Diffuser Llama 3.1-405b

Pour diffuser le modèle Llama 3.1-405b, configurez le transfert de port :

kubectl port-forward svc/jetstream-svc 8000:8000

Le transfert de port vous permet d'accéder au service depuis l'extérieur du cluster. Vous pouvez accéder au déploiement JetStream-Pathways via le service ClusterIP de GKE. Les services ClusterIP ne sont accessibles qu'à partir du cluster.

Interagir avec le modèle

Dans un nouveau terminal, exécutez la commande suivante :

curl --request POST \
--header "Content-type: application/json" \
-s \
localhost:8000/generate \
--data \
'{
    "prompt": "What are the top 5 programming languages",
    "max_tokens": 200
}'

L'exécution de la requête initiale peut prendre plusieurs secondes en raison de l'échauffement du modèle. La sortie devrait ressembler à ce qui suit :

{
    "response": " for web development?\nThe top 5 programming languages for web development are:\n1. **JavaScript**: JavaScript is the most popular language for web development, used by over 90% of websites for client-side scripting. It's also popular for server-side programming with technologies like Node.js.\n2. **HTML/CSS**: HTML (Hypertext Markup Language) and CSS (Cascading Style Sheets) are not programming languages, but are essential for building websites. HTML is used for structuring content, while CSS is used for styling and layout.\n3. **Python**: Python is a popular language for web development, especially with frameworks like Django and Flask. It's known for its simplicity, flexibility, and large community of developers.\n4. **Java**: Java is a popular language for building enterprise-level web applications, especially with frameworks like Spring and Hibernate. It's known for its platform independence, strong security features, and large community of developers.\n5. **PHP**: PHP is a mature language for web"
}

Vous avez effectué les opérations suivantes :

1. Déployé le serveur de modèle JetStream avec MaxText et Pathways sur GKE à l'aide de TPU.
2. Point de contrôle Llama 3.1-405B int8 créé à l'adresse gs://BUCKET_NAME.
3. Le modèle a été diffusé et l'utilisateur a interagi avec lui.

Diffusion désagrégée

Le serving désagrégé est une technique permettant de diffuser des LLM tout en divisant les étapes de préremplissage et de décodage sur différents hôtes. Cette approche optimise l'utilisation des ressources, ce qui peut améliorer le débit et la latence.

• Préremplissage : transmission directe de l'invite d'entrée pour initialiser le cache clé-valeur.

• Décodage : procédure qui génère des jetons de sortie de manière incrémentielle, un jeton par étape et une valeur de cache KV par itération.

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

   export NODE_POOL_NAME=dis-v6e-8
   export NODE_POOL_SIZE=2
   export MACHINE_TYPE=ct6e-standard-4t
   export TOPOLOGY=2x4
   export WORKERS_PER_SLICE=2
   

2. Créez deux pools de nœuds qui utilisent des nœuds v6e-8 :

   for i in $(seq 1 NODE_POOL_SIZE); do
     gcloud container node-pools create NODE_POOL_NAME-${i}-np \
     --project=PROJECT \
     --zone=ZONE \
     --cluster=CLUSTER_NAME \
     --machine-type=MACHINE_TYPE \
     --num-nodes=WORKERS_PER_SLICE \
     --tpu-topology=TOPOLOGY \
     --scopes=cloud-platform \
     --workload-metadata=GCE_METADATA
   done
   

Conversion de point de contrôle

Pour convertir un point de contrôle Meta Llama 2-70B en point de contrôle d'inférence int8 compatible avec MaxText, suivez les étapes décrites dans Conversion de points de contrôle avec Llama2-70B. Sélectionnez Llama2-70B comme modèle lorsque vous acceptez les conditions d'utilisation de Meta. Votre déploiement utilise le point de contrôle avec l'indicateur load_parameters_path.

Remplacez les paramètres suivants dans votre fichier checkpoint-job.yaml :

- --meta_url=META_URL
- --model_name=llama-2
- --model_path=Llama-2-70b-chat
- --output_directory=gs://BUCKET_NAME/maxtext/llama-2-70b

Le point de contrôle sera utilisé dans votre déploiement avec l'indicateur load_parameters_path.

Déployer des parcours JetStream avec un service désagrégé

1. Enregistrez le manifeste suivant sous le nom jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml :

   apiVersion: leaderworkerset.x-k8s.io/v1
   kind: LeaderWorkerSet
   metadata:
     name: jetstream-pathways
     annotations:
       leaderworkerset.sigs.k8s.io/subgroup-exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     replicas: 1
     leaderWorkerTemplate:
       subGroupPolicy:
         subGroupSize: 2
       leaderTemplate:
         metadata:
           labels:
             app: jetstream-pathways
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 2x4
           tolerations:
           - key: "google.com/tpu"
             operator: "Exists"
             effect: "NoSchedule"
           containers:
           - name: pathways-proxy
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-0.5.3
             args:
             - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
             - --server_port=38681
             - --gcs_scratch_location=gs://cloud-pathways-staging/tmp
             - --xla_jf_auto_cross_replica_sharding=false
             - --xla_tpu_enable_windowed_einsum_for_reduce_scatter=false
             - --xla_tpu_enable_windowed_einsum_for_all_gather=false
             - --xla_tpu_prefer_latch_optimized_rhs_layouts=true
             - --xla_tpu_enable_experimental_fusion_cost_model=false
             - --xla_tpu_dot_dot_fusion_duplicated=false
             - --xla_tpu_dot_dot_fusion=true
             - --xla_jf_conv_input_fusion=true
             - --xla_jf_conv_output_fusion=true
             - --xla_tpu_rwb_fusion=false
             - --xla_tpu_copy_fusion_pad_unpad_ratio=0
             - --xla_tpu_licm_size_inflation_ratio=1
             - --xla_tpu_copy_elision_analysis_allowance=150000
             - --xla_tpu_copy_insertion_use_region_analysis_limit=10000
             - --xla_tpu_order_dot_after_layout=true
             - --xla_jf_rematerialization_percent_shared_memory_limit=100
             - --xla_tpu_use_repeated_instance_for_preferred_prefetch_time=true
             - --xla_tpu_enforce_prefetch_fifo_order=false
             - --xla_tpu_prefetch_interval_picker_size_override=6000000
             - --xla_tpu_async_copy_bandwidth_scaling_factor=1
             - --xla_tpu_nd_short_transfer_max_chunks=-1
             - --xla_tpu_enable_aggressive_broadcast_priority_update=true
             - --xla_tpu_alternate_memory_benefit_scaling_factor_for_large_buffers=SQRT
             - --xla_tpu_memory_bound_loop_optimizer_options=enabled:true
             - --xla_tpu_enable_copy_fusion=true
             - --xla_tpu_enable_cross_program_prefetch_freeing=false
             - --xla_tpu_enable_dot_strength_reduction=true
             - --xla_tpu_layout_use_dot_grouping=false
             - --xla_tpu_msa_inefficient_use_to_copy_ratio=0.5
             - --xla_tpu_reduce_loop_fusion_dup_with_unfusable_user=false
             - --xla_tpu_vector_load_fusion_window=1024
             - --xla_tpu_vector_store_fusion_window=256
             - --xla_jf_conv_reshape_fusion=false
             - --xla_tpu_input_conv_multi_users=false
             - --xla_tpu_enable_multi_level_input_dot_dot_fusion=false
             - --xla_tpu_enable_multi_level_output_dot_dot_fusion=false
             - --xla_tpu_dot_dot_fusion_separable_convs_only=false
             - --xla_tpu_enable_multi_level_nested_loop_fusion=true
             - --xla_tpu_nested_dot_fusion=true
             - --xla_tpu_enable_multi_level_nested_dot_fusion=false
             - --xla_jf_enable_multi_output_fusion=true
             - --xla_tpu_use_lp_llo_scheduler_for_dot_dot_fusions=false
             - --xla_tpu_enable_flash_attention=true
             imagePullPolicy: Always
             ports:
             - containerPort: 38681
           - name: pathways-rm
             env:       
             - name: HOST_ADDRESS
               value: "$(LWS_LEADER_ADDRESS)"
             - name: TPU_SKIP_MDS_QUERY
               value: "true"
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             args:
             - --server_port=38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             - --node_type=resource_manager
             - --instance_count=2
             - --instance_type=tpuv6e:2x4
             imagePullPolicy: Always
             ports:
             - containerPort: 38677
           - name: jax-tpu
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-pathways:v0.2.0
             args:
             - MaxText/configs/base.yml
             - tokenizer_path=assets/tokenizer.llama2
             - load_parameters_path=CHECKPOINT_PATH
             - max_prefill_predict_length=1024
             - max_target_length=2048
             - model_name=llama2-70b
             - ici_fsdp_parallelism=1
             - ici_autoregressive_parallelism=1
             - ici_tensor_parallelism=-1
             - scan_layers=false
             - weight_dtype=bfloat16
             - per_device_batch_size=27
             - checkpoint_is_quantized=true 
             - quantization=int8
             - quantize_kvcache=true
             - compute_axis_order=0,2,1,3
             - ar_cache_axis_order=0,2,1,3
             - stack_prefill_result_cache=True
             - inference_server=ExperimentalMaxtextDisaggregatedServer_8
             - inference_benchmark_test=True
             - enable_model_warmup=True
             env:
             - name: LOG_LEVEL
               value: "INFO"
             imagePullPolicy: Always
             securityContext:
               capabilities:
                 add: ["SYS_PTRACE", "NET_ADMIN", "SYS_TIME"]
             ports: 
             - containerPort: 9000
             startupProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 1
               initialDelaySeconds: 240
               failureThreshold: 10000
             livenessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 100
             readinessProbe:
               httpGet:
                 path: /healthcheck
                 port: 8000
                 scheme: HTTP
               periodSeconds: 60
               failureThreshold: 100
           - name: jetstream-http
             image: us-docker.pkg.dev/cloud-tpu-images/inference/jetstream-http:v0.2.3
             imagePullPolicy: Always
             ports:
             - containerPort: 8000
       size: 5
       workerTemplate:
         spec:
           nodeSelector:
             cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
             cloud.google.com/gke-tpu-topology: 2x4
           containers:
           - name: worker
             args:
             - --server_port=38679
             - --resource_manager_address=$(LWS_LEADER_ADDRESS):38677
             - --gcs_scratch_location=PATHWAYS_BUCKET
             image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-0.5.3
             imagePullPolicy: Always
             ports:
             - containerPort: 38679
             resources:
               limits:
                 google.com/tpu: "4"
   --- 
   apiVersion: v1
   kind: Service
   metadata:
     name: jetstream-svc
   spec:
     selector:
       app: jetstream-pathways
     ports:
     - protocol: TCP
       name: jetstream-http
       port: 8000
       targetPort: 8000

2. Définissez la valeur du champ load_parameters_path sur le chemin d'accès de votre point de contrôle qui a été généré lors du processus de conversion des points de contrôle.

   • Pour un point de contrôle bf16, le chemin d'accès doit être semblable à gs://OUTPUT_BUCKET_DIRECTORY/bf16/unscanned/checkpoints/0/items.
   • Pour un point de contrôle int8, il doit être semblable à gs://OUTPUT_BUCKET_DIRECTORY/int8.

   Définissez la valeur du champ gcs_scratch_location sur le bucket Pathways que vous avez créé précédemment.

   perl -pi -e 's|CHECKPOINT_PATH|BUCKET_NAME/maxtext/llama-2-70b/int8|g' jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
   perl -pi -e 's|PATHWAYS_BUCKET|gs://PATHWAYS_BUCKET|g' jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
   

3. Appliquez le fichier manifeste :

   kubectl apply -f jetstream-pathways-disagg-llama-2-70b-2-2x4.yaml
   

   Le serveur de modèle peut mettre un certain temps à restaurer le point de contrôle, en fonction de sa taille. Un modèle de 70 milliards de paramètres peut prendre environ huit minutes pour restaurer le point de contrôle, y compris l'échauffement du modèle. Vous pouvez également observer les journaux pour déterminer le point de préparation en vérifiant le démarrage du serveur de modèle et en diffusant le modèle en configurant le transfert de port pour pouvoir interagir avec le modèle.

Vous avez effectué les opérations suivantes :

1. Déployé le serveur de modèle JetStream avec MaxText et Pathways sur GKE à l'aide de TPU et d'un service désagrégé.
2. Création d'un point de contrôle Llama 2-70B int8 à gs://BUCKET_NAME.
3. Le modèle a été diffusé et l'utilisateur a interagi avec lui.

Résoudre les problèmes

• Si le message Empty reply from server s'affiche, il est possible que le conteneur n'ait pas terminé le téléchargement des 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 le message Connection refused s'affiche, vérifiez que le 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 ne soient facturées sur votre compte Google Cloud , exécutez les commandes suivantes et suivez les instructions :

gcloud container clusters delete CLUSTER_NAME --region=REGION

gcloud iam service-accounts delete jetstream-pathways@PROJECT_ID.iam.gserviceaccount.com

gcloud storage rm --recursive gs://BUCKET_NAME

Étapes suivantes

• Découvrez comment exécuter des modèles Gemma sur GKE et comment exécuter des charges de travail d'IA ou de ML optimisées grâce aux fonctionnalités d'orchestration de plates-formes GKE.
• Apprenez-en plus sur les TPU dans GKE.
• Explorez le dépôt GitHub JetStream.
• Explorez Vertex AI Model Garden.