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.

Ce tutoriel s'adresse aux ingénieurs en machine learning (ML), aux administrateurs et aux opérateurs de plate-forme qui souhaitent utiliser l'orchestration de conteneurs Kubernetes pour gérer des charges de travail d'entraînement, de réglage et d'inférence à grande échelle à l'aide de TPU. Pour en savoir plus sur les rôles courants et les exemples de tâches référencés dans le contenu Google Cloud , consultez la section Rôles utilisateur et tâches courantes de l'utilisateur dans GKE Enterprise.

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.
• Multislice n'est pas compatible avec TPU v3.

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 exécutant une version compatible avec Multislice. Pour connaître les versions compatibles, consultez la section Disponibilité de Multislice 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 un 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_ZONES \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    [--num-nodes=NUM_NODES] \
    [--spot \]
    [--flex-start \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME] \
    [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
    [--placement-type=COMPACT]

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 \
       --location=CONTROL_PLANE_LOCATION
   

   Remplacez les éléments suivants :

   • CLUSTER_NAME : nom du cluster.
   • PROJECT_ID : ID de votre projet
   • CONTROL_PLANE_LOCATION: emplacement Compute Engine du plan de contrôle de votre cluster. Indiquez une région pour les clusters régionaux ou une zone pour les clusters zonaux.

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

   kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=ACCELERATOR_TYPE \
      -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 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.10
                   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.10
                   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 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.

5. (Facultatif) Supprimez la charge de travail JAX :

   kubectl delete -f tpu-multislice.yaml
   

Configurer des paramètres supplémentaires

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

Améliorer les performances réseau avec hostNetwork

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, ajoutez les deux lignes suivantes à la spécification de votre 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.

Si vous utilisez TPU Trillium (v6e) et que vos pods utilisent hostNetworking, installez le DaemonSet suivant pour régler /proc/sys/net/ipv4/tcp_rmem sur le nœud.

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/51bf3dcab6ff658cf62cc32867f96860bf58dfdc/scripts/network-setup/v6e-increase-rmem.yaml

Améliorer les performances réseau sans hostNetwork sur TPU Trillium

Si vous utilisez TPU Trillium et que vos pods ne peuvent pas utiliser hostNetworking, activez la multiréseaux avec le mode netdevice pour obtenir les meilleures performances réseau. La compatibilité de la carte d'interface réseau en mode netdevice avec le multiréseau transmet directement la carte d'interface réseau de la VM au pod, en contournant Kubernetes et GKE Dataplane V2.

Le type de machine ct6e-standard-4t est pris en charge par deux cartes réseau physiques. Kubernetes nécessite un vNIC qui ne peut pas être transmis aux pods. Par conséquent, chaque nœud doit disposer de trois vNIC pour permettre aux pods d'avoir un accès direct à deux vNIC afin d'obtenir les meilleures performances des deux NIC physiques.

Pour activer le mode netdevice pour ct6e-standard-4t, procédez comme suit:

1. Créer deux VPC supplémentaires compatibles avec le mode netdevice
2. Créer un cluster GKE avec des fonctionnalités multiréseaux

3. Configurez deux réseaux netdevice. Par exemple, vous pouvez utiliser les objets GKENetworkParamSet et Network suivants (SECOND_VPC et THIRD_VPC sont les VPC créés à l'étape précédente):

   apiVersion: networking.gke.io/v1
   kind: GKENetworkParamSet
   metadata:
     name: tpu-second
   spec:
     vpc: SECOND_VPC
     vpcSubnet: SECOND_VPC_SUBNET
     deviceMode: NetDevice
   ---
   apiVersion: networking.gke.io/v1
   kind: GKENetworkParamSet
   metadata:
     name: tpu-third
   spec:
     vpc: THIRD_VPC
     vpcSubnet: SECOND_VPC_SUBNET
     deviceMode: NetDevice
   ---
   apiVersion: networking.gke.io/v1
   kind: Network
   metadata:
     name: tpu-second
   spec:
     provider: "GKE"
     type: "Device"
     parametersRef:
       group: networking.gke.io
       kind: GKENetworkParamSet
       name: tpu-second
   ---
   apiVersion: networking.gke.io/v1
   kind: Network
   metadata:
     name: tpu-third
   spec:
     provider: "GKE"
     type: "Device"
     parametersRef:
       group: networking.gke.io
       kind: GKENetworkParamSet
       name: tpu-third
   

4. Connectez vos pods à trois réseaux. Par exemple, vous pouvez utiliser les annotations suivantes dans la spécification de votre pod:

   metadata:
     annotations:
       networking.gke.io/default-interface: 'eth0'
       networking.gke.io/interfaces: |
         [
           {"interfaceName":"eth0","network":"default"},
           {"interfaceName":"eth1","network":"tpu-second"},
           {"interfaceName":"eth2","network":"tpu-third"},
         ]
   

5. Appliquez des sysctls réseau dans le pod, dans le conteneur d'initialisation ou dans le conteneur d'application. Par exemple, vous pouvez ajouter le conteneur d'initialisation suivant à la spécification du pod:

   initContainers:
   - name: "network-optimization-sysctls"
     image: "busybox"
     securityContext:
       privileged: true
     command:
     - bash
     - -c
     - |
       echo 5000 > /proc/sys/net/ipv4/tcp_rto_min_us
       echo 1 > /proc/sys/net/ipv4/tcp_no_metrics_save
       echo 0 > /proc/sys/net/ipv4/tcp_slow_start_after_idle
       echo 131072 > /proc/sys/net/core/optmem_max
       echo "4096 41943040 314572800" > /proc/sys/net/ipv4/tcp_rmem
   

Activer la 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 Afficher les journaux TPU GKE.

Activer des métriques supplémentaires

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 :

Effectuer un nettoyage des ressources

Le moyen le plus simple d'empêcher la facturation est de supprimer le Google Cloud projet que vous avez créé pour ce tutoriel. Vous pouvez également supprimer les différentes ressources.

Supprimer le projet

Delete a Google Cloud project:

gcloud projects delete PROJECT_ID

Supprimer des ressources individuelles

Supprimez le cluster GKE :

```sh
gcloud container clusters delete  CLUSTER_NAME \
   --project=PROJECT_ID  \
   --location=CONTROL_PLANE_LOCATION
```

Étapes suivantes

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