Déployer des charges de travail TPU sur GKE Autopilot

Autopilot

Ce document explique comment accélérer les charges de travail de machine learning (ML) en les déployant à l'aide d'accélérateurs Cloud TPU (TPU) dans des clusters Google Kubernetes Engine (GKE) en mode Autopilot. Avant de lire ce document, assurez-vous de maîtriser les concepts suivants :

Fonctionnement des TPU dans Autopilot

Pour utiliser des TPU dans des charges de travail Autopilot, vous demandez une version de TPU et une topologie compatible pour cette version de TPU dans le fichier manifeste de la charge de travail. Utilisez ensuite les champs Kubernetes resources.requests et resources.limits pour spécifier le nombre de puces TPU que la charge de travail doit utiliser. Lorsque vous déployez la charge de travail, GKE provisionne les nœuds disposant de la configuration TPU demandée et planifie vos pods sur ces nœuds. GKE place chaque charge de travail sur son propre nœud afin que chaque pod puisse accéder à toutes les ressources du nœud avec un risque d'interruption réduit au minimum.

Les TPU dans Autopilot sont compatibles avec les fonctionnalités suivantes :

Planifier la configuration TPU

Avant de demander des TPU, choisissez la configuration souhaitée en fonction des exigences en termes de processeur et de mémoire de la charge de travail. Vous devez décider des éléments suivants :

Version du TPU : version spécifique de Cloud TPU, par exemple la version v5e
Topologie de la version de TPU sélectionnée : disposition et nombre de TPU

La version de TPU et la topologie que vous sélectionnez déterminent si GKE provisionne les nœuds en tant que tranches à hôte unique ou tranches multi-hôte. Dans les tranches à hôte unique, chaque nœud est indépendant des autres nœuds de tranche TPU. Dans les tranches à hôtes multiples, GKE crée un groupe de nœuds ayant des VM interconnectées dans une tranche TPU. Les tranches à hôtes multiples sont atomiques, ce qui signifie que GKE effectue un scaling à la hausse ou à la baisse de l'ensemble du groupe de nœuds interconnectés en une seule unité.

Pour en savoir plus sur les versions de TPU disponibles, les topologies correspondantes, la capacité du processeur et de la mémoire, ainsi que le type de tranche obtenu, consultez la page Choisir une configuration de TPU Autopilot.

Tarifs

Pour en savoir plus sur la tarification, consultez la page Tarifs du mode Autopilot.

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.
Remarque : Pour les installations de gcloud CLI existantes, veillez à définir les propriétés compute/region et compute/zone. En définissant des emplacements par défaut, vous pouvez éviter les erreurs gcloud CLI telles que celle-ci : One of [--zone, --region] must be supplied: Please specify location.

Assurez-vous de disposer d'un cluster Autopilot exécutant la version 1.29.2-gke.1521000 ou ultérieure de GKE.
Pour utiliser des TPU réservés, assurez-vous de disposer d'une réservation de capacité spécifique existante. Pour obtenir des instructions, consultez la section Consommer des ressources zonales réservées.

Vérifier que vous disposez d'un quota de TPU

Pour créer des nœuds de tranche TPU, vous devez disposer d'un quota de TPU, sauf si vous utilisez une réservation de capacité existante. Si vous utilisez des TPU réservés, ignorez cette section.

La création de nœuds de tranche TPU dans GKE nécessite un quota d'API Compute Engine (compute.googleapis.com) et non le quota de l'API Cloud TPU (tpu.googleapis.com). Le nom du quota est différent dans les pods Autopilot standards et dans les pods Spot.

Pour vérifier la limite et l'utilisation actuelle de votre quota d'API Compute Engine pour les TPU, procédez comme suit :

Accédez à la page Quotas de la console Google Cloud.

Accéder à la section "Quotas"
Dans la zone Filtre, procédez comme suit :
1. Sélectionnez la propriété Service, saisissez API Compute Engine, puis appuyez sur Entrée.
2. Sélectionnez la propriété Type et choisissez Quota.
3. Sélectionnez la propriété Nom et saisissez un nom de quota en fonction du type de TPU souhaité, comme suit :
  - TPU v5p (tpu-v5p-slice): puces TPU v5p
  - TPU v5e (tpu-v5-lite-podslice) : puces de tranche de pod TPU v5 Lite
  - TPU v5e (tpu-v5-lite-device) : puces d'appareil TPU v5 Lite
  - TPU v4 (tpu-v4-podslice) : puces de tranche de pod TPU v4
  Pour les pods Spot, sélectionnez le quota "Préemptif" correspondant.
4. Sélectionnez la propriété Dimensions (par exemple, emplacements) et saisissez region: suivi du nom de la région dans laquelle vous prévoyez de créer des TPU dans GKE. Par exemple, saisissez region:us-west4 si vous envisagez de créer des nœuds de tranche TPU dans la zone us-west4-a. Le quota de TPU est régional. Par conséquent, toutes les zones d'une même région consomment le même quota de TPU.

Si aucun quota ne correspond au filtre que vous avez saisi, le projet ne dispose d'aucun quota pour la région souhaitée et vous devez demander une augmentation de quota TPU.

Préparer votre application TPU

Les charges de travail TPU présentent les exigences de préparation suivantes.

Les frameworks tels que JAX, PyTorch et TensorFlow accèdent aux VM des TPU à l'aide de la bibliothèque partagée libtpu. libtpu inclut le compilateur XLA, le logiciel d'environnement d'exécution TPU et le pilote TPU. Chaque version de PyTorch et de JAX nécessite une certaine version de libtpu.so. Pour utiliser des TPU dans GKE, veillez à utiliser les versions suivantes :

Type de TPU	Version de `libtpu.so`
TPU v5e `tpu-v5-lite-podslice` `tpu-v5-lite-device`	Version jax[tpu] recommandée : v0.4.9 ou version ultérieure Version de torchxla[tpuvm] recommandée : v2.1.0 ou version ultérieure
TPU v5p `tpu-v5p-slice`	Version jax[tpu] recommandée : 0.4.19 ou version ultérieure. Version de torchxla[tpuvm] recommandée : il est suggéré d'utiliser une version nocturne du 23 octobre 2023.
TPU v4 `tpu-v4-podslice`	jax[tpu] : v0.4.4 ou version ultérieure torchxla[tpuvm] : v2.0.0 ou version ultérieure

Définissez les variables d'environnement suivantes pour le conteneur qui demande les ressources TPU suivantes :
- TPU_WORKER_ID : entier unique pour chaque pod. Cet ID indique un ID de nœud de calcul unique dans la tranche de TPU. Les valeurs acceptées pour ce champ sont comprises entre zéro et le nombre de pods moins un.
- TPU_WORKER_HOSTNAMES : liste de noms d'hôtes ou d'adresses IP de VM TPU séparés par une virgule qui doivent communiquer entre eux au sein de la tranche. Il doit y avoir un nom d'hôte ou une adresse IP pour chaque VM TPU de la tranche. La liste d'adresses IP ou de noms d'hôte est classée et indexée par zéro par le paramètre TPU_WORKER_ID.

Une fois la préparation de la charge de travail terminée, vous pouvez exécuter un job qui utilise des TPU.

Demander des TPU dans une charge de travail

Cette section explique comment créer un job qui demande des TPU dans Autopilot. Pour toute charge de travail nécessitant des TPU, vous devez spécifier les éléments suivants :

Sélecteurs de nœuds pour la version et la topologie du TPU
Nombre de puces TPU pour un conteneur dans votre charge de travail

Pour obtenir la liste des versions et des topologies de TPU compatibles, ainsi que le nombre correspondant de puces et de nœuds TPU dans une tranche, consultez la page Choisir une configuration de TPU Autopilot.

Remarques concernant les requêtes TPU dans les charges de travail

Un seul conteneur d'un pod peut utiliser des TPU. Le nombre de puces demandé par un conteneur doit être égal au nombre de puces TPU associées à un nœud de la tranche. Par exemple, si vous demandez un TPU v5e (tpu-v5-lite-podslice) avec une topologie 2x4, vous pouvez demander l'un des éléments suivants :

4 puces, qui créent deux nœuds multi-hôtes avec quatre puces TPU chacun
8 puces, qui créent un nœud à hôte unique avec huit puces TPU

Pour optimiser votre rentabilité, il est recommandé de toujours utiliser toutes les puces TPU de la tranche que vous demandez. Si vous demandez une tranche multi-hôte de deux nœuds avec quatre puces TPU chacun, vous devez déployer une charge de travail qui s'exécute sur les deux nœuds et consomme les huit puces de la tranche.

Créer une charge de travail qui demande des TPU

Les étapes suivantes permettent de créer un job qui demande des TPU. Si vous avez des charges de travail qui s'exécutent sur des tranches de TPU multi-hôte, vous devez également créer un service sans adresse IP de cluster qui sélectionne votre charge de travail par son nom. Ce service sans tête IP de cluster permet aux pods de différents nœuds de la tranche multi-hôte de communiquer entre eux en mettant à jour la configuration DNS de Kubernetes de sorte qu'elle pointe vers les pods de la charge de travail.

Enregistrez le manifeste suivant sous le nom tpu-autopilot.yaml :

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS

Remplacez les éléments suivants :

TPU_TYPE : type de TPU à utiliser, tel que tpu-v4-podslice. Il doit s'agir d'une valeur compatible avec GKE.
TOPOLOGY : disposition des puces TPU dans la tranche, par exemple 2x2x4. Il doit s'agir d'une topologie compatible avec le type de TPU sélectionné.
NUMBER_OF_CHIPS : nombre de puces TPU que le conteneur doit utiliser. La valeur doit être identique pour limits et requests.

Déployez le job :
```
kubectl create -f tpu-autopilot.yaml
```

Lorsque vous créez ce job, GKE effectue automatiquement les opérations suivantes :

Il provisionne des nœuds pour exécuter les pods. Selon le type de TPU, la topologie et les demandes de ressources que vous avez spécifiés, ces nœuds sont des tranches à hôte unique ou à plusieurs hôtes.
Il ajoute des rejets aux pods et des tolérances aux nœuds pour éviter que vos autres charges de travail s'exécutent sur les mêmes nœuds que les charges de travail TPU.

Exemple : Afficher le nombre total de puces TPU dans une tranche multi-hôte

La charge de travail suivante renvoie le nombre de puces TPU sur tous les nœuds d'une tranche TPU multi-hôte. Pour créer une tranche multi-hôte, la charge de travail comporte les paramètres suivants :

Version du TPU : TPU v4
Topologie : 2x2x4

Cette version et cette sélection de topologie génèrent une tranche multi-hôte.

Enregistrez le manifeste suivant sous le nom available-chips-multihost.yaml :

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4

Déployez le fichier manifeste :
```
kubectl create -f available-chips-multihost.yaml
```
GKE exécute une tranche de TPU v4 avec quatre VM (tranche TPU multi-hôte). La tranche comporte 16 puces TPU interconnectées.

Vérifiez que le job a créé quatre pods :

kubectl get pods

Le résultat ressemble à ce qui suit :

NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s

Obtenez les journaux de l'un des pods :
```
kubectl logs POD_NAME
```
Remplacez POD_NAME par le nom de l'un des pods créés. Exemple :tpu-job-podslice-0-5cd8r

Le résultat ressemble à ce qui suit :
```
TPU cores: 16
```

Exemple : Afficher les puces TPU dans un seul nœud

La charge de travail suivante est un pod statique qui affiche le nombre de puces TPU associées à un nœud spécifique. Pour créer un nœud à hôte unique, la charge de travail comporte les paramètres suivants :

Version du TPU : TPU v5e
Topologie : 2x4

Cette version et cette sélection de topologie génèrent une tranche à hôte unique.

Enregistrez le manifeste suivant sous le nom available-chips-singlehost.yaml :

apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8

Déployez le fichier manifeste :
```
kubectl create -f available-chips-singlehost.yaml
```
GKE provisionne des nœuds avec huit tranches de TPU à hôte unique qui utilisent les TPU v5e. Chaque nœud TPU possède huit puces TPU (tranche de TPU à hôte unique).
Récupérez les journaux du pod :
```
kubectl logs tpu-job-jax-v5
```
Le résultat ressemble à ce qui suit :
```
Total TPU chips: 8
```

Observabilité et métriques

Tableau de bord

Sur la page Clusters Kubernetes de la console Google Cloud, l'onglet Observabilité affiche les métriques d'observabilité du TPU. Pour en savoir plus, consultez la page Métriques d'observabilité de GKE.

Le tableau de bord TPU n'est renseigné que si les métriques système sont activées dans votre cluster GKE.

Métriques d'exécution

Dans GKE version 1.27.4-gke.900 ou ultérieure, les charges de travail TPU qui utilisent la version JAX 0.4.14 ou ultérieure et spécifient containerPort: 8431 exportent les métriques d'utilisation des TPU en tant que métriques système GKE. Les métriques suivantes sont disponibles dans Cloud Monitoring pour surveiller les performances d'exécution de votre charge de travail TPU :

Cycle d'utilisation : durée exprimée en pourcentage de la dernière période d'échantillonnage (60 secondes) pendant laquelle TensorCore a traité activement les données sur une puce TPU. Un pourcentage plus élevé signifie une meilleure utilisation du TPU.
Mémoire utilisée : quantité de mémoire d'accélérateur allouée, en octets. Cette valeur est échantillonnée toutes les 60 secondes.
Mémoire globale : mémoire totale de l'accélérateur, en octets. Cette valeur est échantillonnée toutes les 60 secondes.

Ces métriques se trouvent dans le schéma de nœud Kubernetes (k8s_node) et de conteneur Kubernetes (k8s_container).

Conteneur Kubernetes :

kubernetes.io/container/accelerator/duty_cycle
kubernetes.io/container/accelerator/memory_used
kubernetes.io/container/accelerator/memory_total

Nœud Kubernetes :

kubernetes.io/node/accelerator/duty_cycle
kubernetes.io/node/accelerator/memory_used
kubernetes.io/node/accelerator/memory_total

Métriques d'hôte

Dans GKE version 1.28.1-gke.1066000 ou ultérieure, les VM dans une tranche TPU exportent les métriques d'utilisation des TPU en tant que métriques système GKE. Les métriques suivantes sont disponibles dans Cloud Monitoring pour surveiller les performances de votre hôte TPU :

Utilisation de TensorCore : pourcentage actuel du TensorCore utilisé. La valeur TensorCore correspond à la somme des unités de multiplication de matrice (MXU) plus l'unité vectorielle. La valeur d'utilisation TensorCore correspond à la division des opérations TensorCore effectuées au cours de la période d'échantillonnage précédente (60 secondes) par le nombre d'opérations TensorCore compatibles sur la même période. Plus la valeur est élevée, plus l'utilisation est optimale.
Utilisation de la bande passante mémoire : pourcentage actuel de bande passante mémoire utilisée par l'accélérateur. Calculé en divisant la bande passante mémoire utilisée sur une période d'échantillonnage (60 secondes) par la bande passante maximale acceptée sur la même période d'échantillonnage.

Ces métriques se trouvent dans le schéma de nœud Kubernetes (k8s_node) et de conteneur Kubernetes (k8s_container).

Conteneur Kubernetes :

kubernetes.io/container/accelerator/tensorcore_utilization
kubernetes.io/container/accelerator/memory_bandwidth_utilization

Nœud Kubernetes :

kubernetes.io/container/node/tensorcore_utilization
kubernetes.io/container/node/memory_bandwidth_utilization

Pour en savoir plus, consultez Métriques de Kubernetes et Métriques système de GKE.

Journalisation

Les journaux émis par les conteneurs s'exécutant sur des nœuds GKE, y compris les VM TPU, sont collectés par l'agent de journalisation GKE, envoyés à Logging et visibles dans Logging.

Recommandations pour les charges de travail TPU dans Autopilot

Les recommandations suivantes peuvent améliorer l'efficacité de vos charges de travail TPU :

Utilisez des pods à durée d'exécution prolongée pendant un délai de grâce allant jusqu'à sept jours avant que GKE n'arrête vos pods pour effectuer des scalings à la baisse ou les mises à niveau de nœuds. Vous pouvez utiliser des intervalles de maintenance et des exclusions avec des pods avec une durée d'exécution étendue pour retarder davantage les mises à niveau automatiques des nœuds.
Utilisez des réservations de capacité pour vous assurer que vos charges de travail reçoivent les TPU demandés sans être placées dans une file d'attente pour disponibilité.