Déployer des TPU Multislice dans GKE

Cette page explique comment déployer des charges de travail dans Google Kubernetes Engine (GKE) à l'aide de la configuration Cloud TPU Multislice pour un entraînement économique et à grande échelle.

Avant de configurer Multislice dans GKE, assurez-vous de connaître les concepts suivants :

1. Présentation de Cloud TPU
2. Architecture du système Cloud TPU
3. À propos des TPU dans GKE

Qu'est-ce que TPU Multislice ?

TPU Multislice est l'organisation architecturale de VM dans une tranche TPU où deux tranches Cloud TPU ou plus communiquent sur le réseau de centre de données. Multislice permet un entraînement full stack à grande échelle et économique, avec un scaling presque linéaire jusqu'à plusieurs dizaines de milliers de puces TPU. Dans une configuration Multislice, GKE déploie une charge de travail Multislice sur plusieurs tranches de TPU. La communication entre les puces TPU d'une tranche s'effectue via des interconnexions entre puces (ICI). La communication entre les tranches s'effectue via le DCN.

Nous vous recommandons d'utiliser Multislice si votre job est trop grand pour tenir sur une seule tranche de TPU.

Disponibilité de Multislice dans GKE

• La version standard est compatible avec Multislice dans les versions 1.27.4-gke.900 et ultérieures.
• Autopilot est compatible avec Multislice dans les versions 1.29.2-gke.1521000 et ultérieures.
• Multislice est compatible avec les frameworks JAX et PyTorch. La version minimale de JAX compatible est la version 2.1.
• Multislice n'est compatible qu'avec les pools de nœuds de tranches TPU multi-hôtes. Par exemple, vous ne pouvez pas utiliser Multislice avec ct4p-hightpu-4t avec une topologie 2x2x1 ou ct5lp-hightpu-4t avec une topologie 2x2, car il s'agit de pools de nœuds de tranche de pod TPU à hôte unique
• Multislice n'est compatible qu'avec l'entraînement synchrone multicontrôleur.
• Les charges de travail Multislice ne peuvent s'exécuter que sur des tranches de TPU partageant le même type, la même taille et la même topologie de TPU.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

• Activez l'API Google Kubernetes Engine.
Activer l'API Google Kubernetes Engine
• Si vous souhaitez utiliser Google Cloud CLI pour cette tâche, installez puis initialisez gcloud CLI. Si vous avez déjà installé gcloud CLI, assurez-vous de disposer de la dernière version en exécutant la commande gcloud components update.

• Créez un cluster standard ou Autopilot qui exécute une version compatible avec Multislice. Pour connaître les versions compatibles, consultez la section Disponibilité de multicouche dans GKE.
• Assurez-vous que votre projet dispose d'un quota suffisant pour Cloud TPU dans GKE.
• Installez JobSet v0.2.3 ou version ultérieure.

Exécuter une charge de travail sur Multislice

Cette section explique comment exécuter une charge de travail sur Multislice. Si vous utilisez le mode GKE Autopilot, passez à la section Exécuter une charge de travail Multislice. Les clusters Autopilot qui exécutent la version 1.29.2-gke.1521000 ou une version ultérieure activent les TPU par défaut.

Préparer un pool de nœuds en mode Standard

Cette section couvre les étapes suivantes :

1. Créer trois pools de nœuds de tranche TPU multi-hôtes.
2. Vérifier l'état du pool de nœuds

Créer le pool de nœuds de tranche TPU

Vous pouvez créer plusieurs pools de nœuds de tranche TPU multi-hôtes. Pour les besoins de ce guide, créez trois pools de nœuds de tranche TPU multi-hôtes pour exécuter une charge de travail Multislice. Vous pouvez créer un pool de nœuds tranche de TPU multi-hôte à l'aide de Google Cloud CLI, de Terraform ou de la console Google Cloud.

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONE \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    --num-nodes=NUM_NODES \
    [--spot \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME]

Vérifier l'état du pool de nœuds

1. Obtenez des identifiants afin de pouvoir utiliser kubectl pour accéder au cluster :

   gcloud container clusters get-credentials CLUSTER_NAME \
       --project=PROJECT_ID
   

   Remplacez les éléments suivants :

   • CLUSTER_NAME : nom du cluster.
   • PROJECT_ID : ID de votre projet

2. Utilisez kubectl dans Cloud Shell pour afficher vos nœuds de tranche TPU :

   kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=TPU_ACCELERATOR \
      -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY
   

   Remplacez les éléments suivants :

   • TPU_ACCELERATOR : type d'accélérateur TPU que vous avez utilisé lors de la création des pools de nœuds. Par exemple, tpu-v4-podslice, tpu-v5-lite-device ou tpu-v5-lite-podslice.
   • TPU_TOPOLOGY : Topologie physique de la tranche de TPU.

   Le résultat ressemble à ce qui suit :

    NAME                                    STATUS   ROLES    AGE    VERSION
    gke-tpu-20ee2cce-5tv6                   Ready    <none>   34h     v1.28.1-gke.1066000
   

Exécuter une charge de travail Multislice

Dans cette section, vous exécutez une charge de travail JAX qui affiche le nombre global de puces TPU dans la tranche TPU, puis se ferme.

Pour exécuter une charge de travail JAX, procédez comme suit :

1. Créez le fichier manifeste tpu-multislice.yaml suivant :

   Autopilot

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-job
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: NUM_SLICES
         template:
           spec:
             parallelism: NUM_NODES
             completions: NUM_NODES
             backoffLimit: 0
             template:
               spec:
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                   cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   - containerPort: 8431
                   command:
                   - bash
                   - -c
                   - |
                     pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                     python -c 'import jax; print("Global device count:", jax.device_count())'
                     sleep 60
                   resources:
                    limits:
                       google.com/tpu: NUM_CHIPS
   

   Standard

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-job
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: NUM_SLICES
         template:
           spec:
             parallelism: NUM_NODES
             completions: NUM_NODES
             backoffLimit: 0
             template:
               spec:
                 hostNetwork: true
                 dnsPolicy: ClusterFirstWithHostNet
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                   cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   - containerPort: 8431
                   securityContext:
                     privileged: true
                   command:
                   - bash
                   - -c
                   - |
                     pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                     python -c 'import jax; print("Global device count:", jax.device_count())'
                     sleep 60
                   resources:
                     limits:
                      google.com/tpu: NUM_CHIPS
   

   Remplacez les éléments suivants :

   • NUM_SLICES: nombre de pools de nœuds de tranche TPU. Dans ce cas, NUM_SLICES est égal à 3.
   • ACCELERATOR_TYPE : type d'accélérateur TPU que vous avez utilisé lors de la création des pools de nœuds. Par exemple, tpu-v4-podslice, tpu-v5-lite-device ou tpu-v5-lite-podslice.
   • TPU_TOPOLOGY : Topologie physique de la tranche de TPU. Par exemple, 4x4x4 ou 2x2 selon la version du TPU.
   • NUM_NODES : Nombre de nœuds dans le pool de nœuds. Doit être égal à zéro ou au produit des valeurs définies dans TPU_TOPOLOGY ({A}x{B}x{C}) divisé par le nombre de puces TPU dans chaque VM.  : pour les TPU multi-hôtes v4, le nombre de puces TPU dans chaque VM est de quatre. Pour les TPU multi-hôtes v5e, le nombre de puces TPU dans chaque VM est de un, quatre ou huit. Par conséquent, si votre TPU_TOPOLOGY est 2x4x4 (TPU v4 avec quatre puces TPU dans chaque VM), la valeur de NUM_NODES est 32/4, ce qui correspond à 8.
   • NUM_CHIPS : pour les TPU multi-hôtes v4, le nombre de puces TPU dans chaque VM est de quatre. Pour les TPU multi-hôtes v5e, le nombre de puces TPU dans chaque VM est de un, quatre ou huit. Pour en savoir plus, consultez la section puces TPU sur la VM dans une tranche TPU.

   Dans le fichier manifeste :

   • JobSet est un service headless portant le même nom que le nom JobSet. Dans le cas présent, il s'agit de multislice-job.
   • L'annotation alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool configure l'affinité des pods pour s'assurer que tous les pods sont programmés sur la même tranche.
   • La valeur maxRestarts: 4 indique le nombre maximal de redémarrages de GKE sur une ressource JobSet en cas d'échec d'un job enfant. Si le redémarrage de JobSet atteint la valeur maximale définie, la ressource JobSet est marquée comme ayant échoué.
   • Les champs parallelism et completions correspondent au nombre de nœuds dans chaque pool de nœuds.
   • Le paramètre backoff est égal à 0, car Multislice n'accepte que l'entraînement synchrone multicontrôleurs. Doit être défini sur 0. Faire échouer le job en cas de défaillance d'un pod.
   • Les valeurs de la section d'affinité permettent de s'assurer qu'une seule charge de travail TPU Multislice s'exécute dans un groupe de segments multiples.
   • containerPort: 8080 est le port du coordinateur MXLA.
   • containerPort: 8431 est le port permettant d'exporter les métriques d'utilisation du TPU.
   • L'élément securityContext: privileged: true indique que le mode privilégié est activé pour les nœuds pour l'accès aux TPU. Les nœuds de GKE version 1.28 ou ultérieure n'ont pas besoin d'activer le mode privilégié pour accéder aux TPU. Pour en savoir plus, consultez la page Exécuter des conteneurs sans mode privilégié.

2. Appliquez le fichier manifeste :

   kubectl apply -f tpu-multislice.yaml
   

3. Confirmez que la charge de travail est acceptée :

   kubectl get jobsets
   

   Le résultat ressemble à ce qui suit :

   NAME            RESTARTS   COMPLETED   AGE
   multislice-job                         3s
   

4. Surveillez l'état des pods provisionnés :

   kubectl get pods
   

   Le résultat ressemble à ce qui suit :

    NAME                                READY   STATUS      RESTARTS   AGE
    multislice-job-slice-0-0-wzq9t      0/1     Completed   0          2m31s
    multislice-job-slice-0-1-zf4dp      0/1     Completed   0          2m30s
    multislice-job-slice-1-0-hbfn5      0/1     Completed   0          2m31s
    multislice-job-slice-1-1-45fgl      0/1     Completed   0          2m30s
    multislice-job-slice-2-0-wjbp4      0/1     Completed   0          2m30s
    multislice-job-slice-2-1-lwnvs      0/1     Completed   0          2m30s
   

JobSet multislice-job planifie, crée, puis exécute les pods jusqu'à la fin. Les noms des pods sont au format <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>. Le préfixe jobsetName détermine le JobJob auquel appartient le pod.

Configurations supplémentaires

Les sections suivantes décrivent les configurations supplémentaires que vous pouvez appliquer à votre objet Multislice.

Activer hostNetwork sur vos pods GKE Standard

Pour améliorer les performances réseau entre les tranches de TPU, nous vous recommandons d'activer hostNetworking. Utilisez hostNetwork: true dans la spécification de pod pour ignorer toute la pile réseau Kubernetes et laisser vos pods Kubernetes utiliser directement le réseau hôte pour une communication de VM à VM.

Pour activer hostNetworking, supprimez les deux lignes suivantes de la spécification de pod :

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

Pour continuer à utiliser podHostnames pour la découverte de nœuds de calcul avec hostNetwork, définissez dnsPolicy: ClusterFirstWithHostNet. Cela est important lorsque vous exécutez des tâches de reprise automatique et que vous devez disposer des mêmes noms pour actualiser les mêmes points de contrôle.

Journalisation

Les journaux émis par les conteneurs s'exécutant sur des nœuds GKE, y compris les nœuds de tranche TPU, sont visibles dans l'explorateur de journaux, si vous avez activé la journalisation système GKE dans votre cluster.

Vous pouvez afficher vos journaux à partir de GKE à l'aide de l'explorateur de journaux avec le filtre suivant pour afficher les journaux de conteneur de votre charge de travail :

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME

Utilisez le filtre suivant pour les tranches et les nœuds de calcul TPU :

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>

Pour en savoir plus, consultez la section Afficher les journaux GKE TPU.

Observabilité et métriques

En plus des métriques TPU générales, il existe quatre métriques d'exécution TPU multislice spécifiques. Ces métriques sont disponibles dans GKE 1.29.1-gke.1016000 ou version ultérieure. La charge de travail TPU doit utiliser la version JAX 0.4.24

Voici les métriques multislice disponibles :

• Latences de transfert DCN (réseau de centres de données) : distribution des latences de transfert réseau pour le trafic multislice.
• Latences collectives : répartition de la latence collective de bout en bout pour le trafic multislice.
• Latences de transfert hôte à appareil : répartition de la latence de transfert hôte à appareil pour chaque fragment de données pour le trafic multislice.
• Latences de transfert d'appareil à hôte : répartition de la latence de transfert entre l'appareil et l'hôte, pour chaque fragment de données pour le trafic multislice.

Ces métriques se trouvent dans le schéma du conteneur Kubernetes (k8s_container) :

• kubernetes.io/container/multislice/network/dcn_transfer_latencies
• kubernetes.io/container/multislice/network/collective_end_to_end_latencies
• kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies
• kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies

Tranche de pod TPU ou Multislice

Le tableau suivant différencie l'organisation architecturale d'une tranche de TPU et Multislice :

Étapes suivantes

• Découvrez comment orchestrer des charges de travail Multislice avec des tranches de TPU.