Déployer des charges de travail TPU dans GKE

Standard

Cette page explique comment demander et déployer des charges de travail qui utilisent des accélérateurs Cloud TPU (TPU) dans Google Kubernetes Engine (GKE).

Avant de configurer et déployer des charges de travail TPU dans GKE, vous devez connaître les concepts suivants :

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.

Disponibilité des TPU dans GKE

Utilisez GKE pour créer et gérer des pools de nœuds à l'aide de TPU. Vous pouvez utiliser ces accélérateurs sur mesure pour effectuer l'entraînement, le réglage et l'inférence de modèles d'IA à grande échelle.

Consultez la liste des versions de TPU acceptées dans GKE.

Planifier la configuration TPU

Planifiez la configuration de votre TPU en fonction de votre modèle de machine learning et de la quantité de mémoire requise. Voici les étapes pertinentes lors de la planification de votre configuration TPU :

Garantir un quota suffisant

Pour vérifier votre limite de quota et votre utilisation actuelle de TPU, suivez la procédure ci-dessous.

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

Accéder à la section "Quotas"

Dans la zone Filtre, procédez comme suit :

Sélectionnez la propriété Service, saisissez API Compute Engine, puis appuyez sur Entrée.

Sélectionnez la propriété Quota, puis saisissez le nom du quota en fonction de la version du TPU et du type de machine. Par exemple, si vous envisagez de créer des nœuds TPU v5e à la demande dont le type de machine commence par ct5lp-, saisissez TPU v5 Lite PodSlice chips.

Version du TPU	Type de machine commençant par	Nom du quota pour les instances à la demande	Nom du quota pour les instances de VM Spot¹	Nom du quota pour les instances réservées²
TPU v4	`ct4p-`	`TPU v4 PodSlice chips`	`Preemptible TPU v4 PodSlice chips`	`Committed TPU v4 PodSlice chips`
TPU v5e	`ct5l-`	`TPU v5 Lite Device chips`	`Preemptible TPU v5 Lite Device chips`	`Committed TPU v5 Lite Device chips`
TPU v5e	`ct5lp-`	`TPU v5 Lite PodSlice chips`	`Preemptible TPU v5 Lite PodSlice chips`	`Committed TPU v5 Lite PodSlice chips`

Vous pouvez éventuellement appliquer des filtres plus avancés pour affiner les résultats. Pour ce faire, sélectionnez la propriété Dimensions (emplacements, par exemple), ajoutez le nom de la région dans laquelle vous souhaitez créer des TPU dans GKE, puis appuyez sur Entrée. Par exemple, saisissez region:us-west4 si vous envisagez de créer des nœuds TPU v5e dans la zone us-west4-a.

Si aucun quota ne correspond au filtre que vous avez saisi, le projet ne dispose d'aucun quota et vous devez demander une augmentation de quota TPU.

Lorsque vous créez un pool de nœuds TPU, utilisez l'option --spot pour créer une instance de VM Spot. ↩
Lorsque vous créez un pool de nœuds TPU, utilisez l'option --reservation pour créer une instance réservée. Les réservations TPU sont disponibles lors de la souscription d'un engagement. ↩

Demander une augmentation du quota de TPU

Si vous avez besoin d'un quota supplémentaire pour créer des nœuds TPU dans GKE, demandez une augmentation de votre quota de TPU en contactant votre responsable de compte Google Cloud. Les TPU provisionnés avec GKE utilisent un quota d'allocation de l'API Compute Engine. Le quota demandé pour le quota d'API Cloud TPU ne s'applique pas aux TPU dans GKE.

Créer un cluster

Créez un cluster GKE en mode standard dans une région avec des TPU disponibles. Nous vous recommandons d'utiliser des clusters régionaux, qui offrent une haute disponibilité du plan de contrôle Kubernetes. Vous pouvez utiliser Google Cloud CLI ou la console Google Cloud.

gcloud container clusters create CLUSTER_NAME \
  --location LOCATION

Remplacez les éléments suivants :

CLUSTER_NAME : nom du nouveau cluster
LOCATION : région dans laquelle votre capacité de TPU est disponible.
VERSION : numéro de version de GKE. GKE est compatible avec les TPU v4 dans les versions 1.26.1.gke-1500 et ultérieures, et les TPU v5e dans les versions 1.27.2-gke.2100 et ultérieures. Pour en savoir plus, consultez la section Disponibilité des TPU dans GKE.

Créer un pool de nœuds

Tranche de TPU à hôte unique

Vous pouvez créer un pool de nœuds de tranche de TPU à hôte unique à l'aide de Google Cloud CLI, de Terraform ou de la console Google Cloud.

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    [--num-nodes=NUM_NODES \]
    [--spot \]
    [--enable-autoscaling \]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME \]
    [--total-min-nodes TOTAL_MIN_NODES \]
    [--total-max-nodes TOTAL_MAX_NODES \]
    [--location-policy=ANY]

Remplacez les éléments suivants :

POOL_NAME : nom du nouveau pool de nœuds.
LOCATION : nom de la zone basé sur la version de TPU que vous souhaitez utiliser :
- Pour les TPU v4, utilisez us-central2-b.
- Pour les types de machines TPU v5e commençant par ct5l-, utilisez us-central1-a ou europe-west4-b.
- Pour les types de machines TPU v5e commençant par ct5lp-, utilisez us-west4-a, us-east1-c ou us-east5-b. Pour en savoir plus, consultez la section Sélectionner une version et une topologie de TPU.
CLUSTER_NAME : nom du cluster.
NODE_ZONE : liste d'une ou de plusieurs zones, séparées par une virgule, dans lesquelles GKE crée le pool de nœuds.
MACHINE_TYPE : type de machine à utiliser pour les nœuds. Pour en savoir plus sur les types de machines compatibles avec les TPU, utilisez le tableau de la section Mappage de la configuration de TPU.

Vous pouvez également utiliser les options suivantes :

NUM_NODES : nombre initial de nœuds dans le pool de nœuds de chaque zone. Si vous omettez cette option, la valeur par défaut est 3. Si l'autoscaling est activé pour le pool de nœuds à l'aide de l'option --enable-autoscaling, nous vous recommandons de définir NUM_NODES sur 0, car l'autoscaler provisionne des nœuds supplémentaires dès que vos charges de travail en ont besoin.
RESERVATION_NAME : nom de la réservation utilisée par GKE lors de la création du pool de nœuds. Si vous omettez cette option, GKE utilise les TPU disponibles. Pour en savoir plus sur les réservations TPU, consultez la page Réservations TPU.
--enable-autoscaling : Ajouter un pool de nœuds avec l'autoscaling activé
- TOTAL_MIN_NODES : nombre minimal de nœuds dans le pool de nœuds. Omettez ce champ, sauf si l'autoscaling est également spécifié.
- TOTAL_MAX_NODES : nombre maximal de nœuds dans le pool de nœuds. Omettez ce champ, sauf si l'autoscaling est également spécifié.
--spot : définit le pool de nœuds de sorte qu'il utilise des VM Spot pour les nœuds du pool de nœuds. Ce paramètre ne peut pas être modifié après la création du pool de nœuds.

Terraform

Assurez-vous d'utiliser la version 4.84.0 ou ultérieure du fournisseur google.
Ajoutez le bloc suivant à votre configuration Terraform :

resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES
  autoscaling {
    total_min_node_count = TOTAL_MIN_NODES
    total_max_node_count = TOTAL_MAX_NODES
    location_policy      = "ANY"
  }

  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
  }
}

Remplacez les éléments suivants :

NODE_POOL_RESOURCE_NAME : nom de la ressource du pool de nœuds dans le modèle Terraform.
PROJECT_ID : ID de votre projet
CLUSTER_NAME : nom du cluster existant.
POOL_NAME : nom du pool de nœuds à créer
CLUSTER_LOCATION : zone de calcul du cluster. Spécifiez la région dans laquelle la version de TPU est disponible. Pour en savoir plus, consultez la page Sélectionner une version et une topologie de TPU.
NODE_ZONES : liste d'une ou de plusieurs zones, séparées par une virgule, dans lesquelles GKE crée le pool de nœuds.
NUM_NODES : nombre initial de nœuds dans le pool de nœuds de chacune des zones du pool. Si aucune valeur n'est spécifiée, la valeur par défaut est 3. Si l'autoscaling est activé pour le pool de nœuds à l'aide du modèle d'autoscaling, nous vous recommandons de définir NUM_NODES sur 0, car GKE provisionne des nœuds TPU supplémentaires dès que votre charge de travail les requiert.
MACHINE_TYPE : type de machine TPU à utiliser. Pour afficher les types de machines TPU compatibles, utilisez le tableau de la section Mappage de la configuration de TPU.

Vous pouvez également utiliser les variables suivantes :

autoscaling : Ajouter un pool de nœuds avec l'autoscaling activé. Pour la tranche TPU à hôte unique, GKE effectue un scaling entre les valeurs TOTAL_MIN_NODES et TOTAL_MAX_NODES.
- TOTAL_MIN_NODES : nombre minimal de nœuds dans le pool de nœuds. Ce champ est facultatif, sauf si l'autoscaling est également spécifié.
- TOTAL_MAX_NODES : nombre maximal de nœuds dans le pool de nœuds. Ce champ est facultatif, sauf si l'autoscaling est également spécifié.
RESERVATION_NAME : si vous utilisez la réservation TPU, il s'agit de la liste des libellés des ressources de réservation à utiliser lors de la création du pool de nœuds. Pour en savoir plus sur la spécification de RESERVATION_LABEL_VALUES dans le champ reservation_affinity, consultez la page Fournisseur Terraform.
spot : définit le pool de nœuds de sorte qu'il utilise des VM Spot pour les nœuds TPU. Ce paramètre ne peut pas être modifié après la création du pool de nœuds. Pour en savoir plus, consultez la page VM Spot.

Console

Pour créer un pool de nœuds à l'aide de TPU, exécutez la commande suivante :

Accédez à la page Google Kubernetes Engine dans Google Cloud Console.

Accéder à Google Kubernetes Engine
Dans la liste des clusters, cliquez sur le nom du cluster que vous souhaitez modifier.
Cliquez sur Ajouter un pool de nœuds.
Dans la section Détails du pool de nœuds, cochez la case Spécifier les emplacements de nœud.
Sélectionnez la zone en fonction de la version de TPU que vous souhaitez utiliser :
- Pour les TPU v4, utilisez us-central2-b.
- Pour les types de machines TPU v5e commençant par ct5l-, utilisez us-central1-a ou europe-west4-b.
- Pour les types de machines TPU v5e commençant par ct5lp-, utilisez us-west4-a, us-east1-c ou us-east5-b.
Dans le volet de navigation, cliquez sur Nœuds.
Dans la section Configuration de la machine, sélectionnez TPU.
Dans le menu déroulant Série, sélectionnez l'une des options suivantes :
- CT4P : TPU v4
- CT5LP : TPU v5e
Dans le menu déroulant Type de machine, sélectionnez le nom de la machine à utiliser pour les nœuds. Utilisez le tableau Mappage des configurations TPU pour apprendre à définir le type de machine et la topologie TPU qui créent un pool de nœuds TPU à hôte unique.
Dans le menu déroulant Topologie TPU, sélectionnez la topologie physique de la tranche de TPU.
Dans la boîte de dialogue Modifications requises, cliquez sur Apporter des modifications.
Assurez-vous d'avoir défini le champ Type de disque de démarrage sur Disque persistant standard ou Disque persistant SSD.
Vous pouvez éventuellement cocher la case Activer les nœuds sur les VM Spot pour utiliser des VM Spot pour les nœuds du pool de nœuds.
Cliquez sur Créer.

Tranche de TPU multi-hôte

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

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]

Remplacez les éléments suivants :

POOL_NAME : nom du nouveau pool de nœuds.
LOCATION : nom de la zone basé sur la version de TPU que vous souhaitez utiliser :
- Pour les TPU v4, utilisez us-central2-b.
- Pour les types de machines TPU v5e commençant par ct5lp-, utilisez us-west4-a, us-east1-c ou us-east5-b. Les types de machines TPU v5e commençant par ct5l- ne sont jamais à hôtes multiples. Pour en savoir plus, consultez la section Sélectionner une version et une topologie de TPU.
CLUSTER_NAME : nom du cluster.
NODE_ZONE : liste d'une ou de plusieurs zones, séparées par une virgule, dans lesquelles GKE crée le pool de nœuds.
MACHINE_TYPE : type de machine à utiliser pour les nœuds. Pour en savoir plus sur les types de machines disponibles, consultez la section Mappage de la configuration TPU.
TPU_TOPOLOGY : Topologie physique de la tranche de TPU. Le format de la topologie dépend de la version du TPU comme suit :
- TPU v4 : définissez la topologie à trois tuples ({A}x{B}x{C}), par exemple 4x4x4.
- TPU v5e : définissez la topologie à deux tuples ({A}x{B}), par exemple 2x2.
NUM_NODES : Nombre de nœuds dans le pool de nœuds. Elle doit être égale à zéro ou au produit des valeurs définies dans TPU_TOPOLOGY ({A}x{B}x{C}) divisée par le nombre de puces dans chaque VM. Pour les TPU multi-hôtes v4 et v5e, le nombre de puces dans chaque VM est de quatre. Par conséquent, si votre TPU_TOPOLOGY est de 2x4x4 (TPU v4 avec quatre puces dans chaque VM), la valeur de NUM_NODES est 32/4, ce qui correspond à 8.

Vous pouvez également utiliser les options suivantes :

RESERVATION_NAME : nom de la réservation utilisée par GKE lors de la création du pool de nœuds. Si vous omettez cette option, GKE utilise les pools de nœuds TPU disponibles. Pour en savoir plus sur les réservations TPU, consultez la page Réservations TPU.
--spot : définit le pool de nœuds de sorte qu'il utilise des VM Spot pour les nœuds TPU. Ce paramètre ne peut pas être modifié après la création du pool de nœuds. Pour en savoir plus, consultez la page VM Spot.
--enable-autoscaling : Ajouter un pool de nœuds avec l'autoscaling activé. Lorsque GKE met à l'échelle un pool de nœuds de tranche TPU multi-hôte, il effectue un scaling de façon atomique du pool de nœuds de zéro à la taille maximale.
- MAX_NODES : taille maximale du pool de nœuds. Elle doit être égale au produit des valeurs définies dans TPU_TOPOLOGY ({A}x{B}x{C}) divisée par le nombre de puces dans chaque VM.

Terraform

Assurez-vous d'utiliser la version 4.84.0 ou ultérieure du fournisseur google.
Ajoutez le bloc suivant à votre configuration Terraform :
```
resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES

  autoscaling {
    max_node_count = MAX_NODES
    location_policy      = "ANY"
  }
  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = TPU_TOPOLOGY
  }
}
```
Remplacez les éléments suivants :
- NODE_POOL_RESOURCE_NAME : nom de la ressource de pool de nœuds dans le modèle Terraform.
- PROJECT_ID : ID de votre projet
- CLUSTER_NAME : nom du cluster existant auquel ajouter le pool de nœuds.
- POOL_NAME : nom du pool de nœuds à créer
- CLUSTER_LOCATION : emplacement de calcul du cluster. Nous vous recommandons de disposer d'un cluster régional pour améliorer la fiabilité du plan de contrôle Kubernetes. Vous pouvez également utiliser un cluster zonal. Pour en savoir plus, consultez la page Sélectionner une version et une topologie de TPU.
- NODE_ZONES : liste d'une ou de plusieurs zones, séparées par une virgule, dans lesquelles GKE crée le pool de nœuds.
- NUM_NODES : Nombre de nœuds dans le pool de nœuds. Doit être égal à zéro ou au produit du nombre de puces TPU divisé par quatre, car dans les tranches de TPU à hôtes multiples, chaque nœud TPU comporte quatre puces. Par exemple, si TPU_TOPOLOGY est défini sur 4x8, il y a 32 puces, ce qui signifie que NUM_NODES doit être égal à 8. Pour en savoir plus sur les topologies TPU, utilisez le tableau de la section Mappage de la configuration TPU.
- TPU_TOPOLOGY : indique la topologie physique souhaitée pour la tranche de TPU. Le format de la topologie dépend de la version de TPU que vous utilisez :
  - Pour les TPU v4 : définissez la topologie à trois tuples ({A}x{B}x{C}), par exemple 4x4x4.
  - Pour les TPU v5e, définissez la topologie en deux tuples ({A}x{B}), par exemple 2x2.
Vous pouvez également utiliser les variables suivantes :
- RESERVATION_NAME : si vous utilisez la réservation TPU, il s'agit de la liste des libellés des ressources de réservation à utiliser lors de la création du pool de nœuds. Pour en savoir plus sur la spécification de RESERVATION_LABEL_VALUES dans le champ reservation_affinity, consultez la page Fournisseur Terraform.
- autoscaling : Ajouter un pool de nœuds avec l'autoscaling activé. Lorsque GKE met à l'échelle un pool de nœuds de tranche TPU multi-hôte, il effectue un scaling de façon atomique du pool de nœuds de zéro à la taille maximale.
  - MAX_NODES: taille maximale du pool de nœuds. Elle doit être égale au produit des valeurs définies dans TPU_TOPOLOGY ({A}x{B}x{C}) divisée par le nombre de puces dans chaque VM.
- spot : permet au pool de nœuds d'utiliser des VM Spot pour les nœuds TPU. Ce paramètre ne peut pas être modifié après la création du pool de nœuds. Pour en savoir plus, consultez la page VM Spot.

Console

Pour créer un pool de nœuds à l'aide de TPU, exécutez la commande suivante :

Accédez à la page Google Kubernetes Engine dans Google Cloud Console.

Accéder à Google Kubernetes Engine
Dans la liste des clusters, cliquez sur le nom du cluster que vous souhaitez modifier.
Cliquez sur Ajouter un pool de nœuds.
Dans la section Détails du pool de nœuds, cochez la case Spécifier les emplacements de nœud.
Sélectionnez la zone en fonction de la version de TPU que vous souhaitez utiliser :
- Pour les TPU v4, utilisez us-central2-b.
- Pour les TPU v5e, utilisez us-west4-a ou us-east1-c.
Dans le volet de navigation, cliquez sur Nœuds.
Dans la section Configuration de la machine, sélectionnez TPU.
Dans le menu déroulant Série, sélectionnez l'une des options suivantes :
- CT4P : pour les TPU v4.
- CT5LP : pour les TPU v5e
Dans le menu déroulant Type de machine, sélectionnez le nom de la machine à utiliser pour les nœuds. Utilisez le tableau Mappage des configurations TPU pour apprendre à définir le type de machine et la topologie TPU qui créent un pool de nœuds TPU multi-hôtes.
Dans le menu déroulant Topologie TPU, sélectionnez la topologie physique de la tranche de TPU.
Dans la boîte de dialogue Modifications requises, cliquez sur Apporter des modifications.
Assurez-vous d'avoir défini le champ Type de disque de démarrage sur Disque persistant standard ou Disque persistant SSD.
Vous pouvez éventuellement cocher la case Activer les nœuds sur les VM Spot afin d'utiliser des VM Spot pour les nœuds du pool de nœuds.
Cliquez sur Créer.

Exécuter votre charge de travail sur des nœuds TPU

Préparation des charges de travail

Après avoir créé votre cluster et vos pools de nœuds, vous pouvez configurer vos charges de travail. Au préalable, vous devez suivre les étapes de préparation des charges de travail 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, l'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 :

Version du TPU	Version de l'application `libtpu.so`
TPU v4	Version jax[tpu] recommandée : 0.4.4 ou version ultérieure. Version de torchxla[tpuvm] recommandée : v2.0.0 ou version ultérieure.
TPU v5e	Version jax[tpu] recommandée : v0.4.9 ou version ultérieure. Version de torchxla[tpuvm] recommandée : il est suggéré d'utiliser une version nocturne du 17/07 (gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm_20230717).

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.
GKE injecte automatiquement ces variables d'environnement en utilisant un webhook en mutation lorsqu'une tâche est créée avec completionMode: Indexed, subdomain, parallelism > 1, et en demandant des propriétés google.com/tpu. GKE ajoute un service sans interface graphique afin que les enregistrements DNS soient ajoutés pour les pods qui sauvegardent le service.
Dans le fichier manifeste de votre charge de travail, ajoutez des sélecteurs de nœuds Kubernetes pour vous assurer que GKE planifie votre charge de travail TPU sur le type de machine et la topologie de TPU que vous avez définis :
```
nodeSelector:
  cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
  cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
```

Remplacez les éléments suivants :

TPU_ACCELERATOR : nom de l'accélérateur TPU :
- Pour les TPU v4, utilisez tpu-v4-podslice.
- Pour les types de machines TPU v5e commençant par ct5l-, utilisez tpu-v5-lite-device.
- Pour les types de machines TPU v5e commençant par ct5lp-, utilisez tpu-v5-lite-podslice.
TPU_TOPOLOGY : topologie physique de la tranche de pod TPU. Le format de la topologie dépend de la version du TPU :
- TPU v4 : définissez la topologie à trois tuples ({A}x{B}x{C}), par exemple 4x4x4.
- TPU v5e : définissez la topologie à deux tuples ({A}x{B}), par exemple 2x2.

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

Les sections suivantes montrent des exemples d'exécution d'un job qui effectue des calculs simples avec des TPU.

Exemple 1 : Exécuter une charge de travail qui affiche le nombre de puces TPU disponibles dans un pool de nœuds TPU

Cet exemple inclut la configuration suivante :

Version du TPU : v4 (tpu-v4-podslice)
Topologie : 2x2x4
Type de pool de nœuds : tranche de TPU multi-hôte

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

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-podslice
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-podslice
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.
        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("TPU cores:", jax.device_count())'
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Chaque nœud TPU possède les libellés de nœud suivants :

cloud.google.com/gke-accelerator-type: tpu-v4-podslice
cloud.google.com/gke-tpu-topology: 2x2x4

Appliquez le fichier manifeste :
```
kubectl apply -f tpu-job.yaml
```
GKE exécute une tranche de TPU v4 avec quatre VM TPU (tranche TPU multi-hôte). Le pool de nœuds comporte 16 puces interconnectées.

Vérifiez que la tâche 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. Par exemple, tpu-job-podslice-0-5cd8r.

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

Exemple 2 : exécuter une charge de travail qui affiche le nombre de puces TPU disponibles dans la VM TPU

Topologie : 2x4
Version du TPU : v5e (tpu-v5-lite-podslice)
Type de pool de nœuds : tranche de TPU à hôte unique

Créez le fichier manifeste suivant en tant que tpu-pod.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.
    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("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8

Ce fichier manifeste comprend les libellés de nœud suivants :

cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
cloud.google.com/gke-tpu-topology: 2x4

GKE provisionne un pool de nœuds avec huit tranches de TPU à hôte unique qui utilisent les TPU v5e. Chaque VM TPU possède huit puces (tranche de TPU à hôte unique).

Effectuer un nettoyage

Pour éviter que les ressources utilisées dans ce guide soient facturées sur votre compte Google Cloud, envisagez de supprimer les pools de nœuds TPU qui ne sont plus associés à des charges de travail planifiées. Si les charges de travail en cours d'exécution doivent être arrêtées correctement, utilisez kubectl drain pour nettoyer les charges de travail avant de supprimer le nœud.

Supprimez un pool de nœuds TPU :
```
gcloud container node-pools delete POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME
```
Remplacez les éléments suivants :
- POOL_NAME : nom du pool de nœuds.
- CLUSTER_NAME : nom du cluster.
- LOCATION : emplacement de calcul du cluster.

Configurations supplémentaires

Les sections suivantes décrivent les configurations supplémentaires que vous pouvez appliquer à vos charges de travail TPU.

Multicouche

Vous pouvez agréger des tranches plus petites dans une tranche multicouche pour gérer des charges de travail d'entraînement plus volumineuses. Pour en savoir plus, consultez la section TPU multicouche dans GKE.

Migrer votre réservation TPU

Si vous disposez de réservations TPU existantes, vous devez d'abord migrer votre réservation TPU vers un nouveau système de réservation basé sur Compute Engine. Vous pouvez également créer un système de réservation basé sur Compute Engine dans lequel aucune migration n'est nécessaire. Pour savoir comment migrer vos réservations TPU, consultez la page Réservation TPU.

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.

Utiliser le provisionnement automatique des nœuds GKE

Vous pouvez configurer GKE pour créer et supprimer automatiquement des pools de nœuds afin de répondre aux exigences de ressources de vos charges de travail TPU. Pour en savoir plus, consultez la page Configurer Cloud TPU.

Réparation automatique des nœuds TPU

Si un nœud TPU dans un pool de nœuds TPU multi-hôte n'est pas opérationnel, l'intégralité du pool de nœuds est recréé. Les conditions qui entraînent des nœuds TPU non opérationnels sont les suivantes :

Tout nœud TPU avec des conditions de nœuds communes.
Tout nœud TPU dont le nombre de TPU non allouable est supérieur à zéro.
Toute instance de VM TPU arrêtée (en raison de la préemption) ou terminée.
Maintenance des nœuds : si un nœud TPU (VM) d'un pool de tranches de TPU multi-hôte est défaillant pour la maintenance de l'hôte, GKE recrée l'ensemble de la tranche de TPU.

Vous pouvez consulter l'état de la réparation (y compris le motif de l'échec) dans l'historique des opérations. Si l'échec est dû à un quota insuffisant, contactez votre responsable de compte Google Cloud pour augmenter le quota correspondant.

Observabilité et métriques

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 sont situées dans le schéma du nœud Kubernetes (k8s_node) et du 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 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 matricielles (MXU) et de 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.
Une valeur plus élevée signifie une meilleure utilisation.
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 sont situées dans le schéma du nœud Kubernetes (k8s_node) et du 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.

Exécuter des conteneurs sans mode privilégié

Si le nœud TPU exécute une version antérieure à 1.28, lisez la section suivante :

Un conteneur s'exécutant sur une VM TPU doit avoir accès à des limites plus élevées sur la mémoire verrouillée, afin que le pilote puisse communiquer avec les puces TPU via l'accès direct à la mémoire (DMA). Pour activer cette fonctionnalité, vous devez configurer une valeur ulimit plus élevée. Si vous souhaitez réduire le champ d'application des autorisations sur votre conteneur, procédez comme suit :

Modifiez le fichier securityContext afin d'inclure les champs suivants :
```
securityContext:
  capabilities:
    add: ["SYS_RESOURCE"]
```
Augmentez ulimit en exécutant la commande suivante dans le conteneur avant de configurer vos charges de travail pour qu'elles utilisent des ressources TPU :
```
ulimit -l 68719476736
```

Remarque : Pour les TPU v5e, l'exécution de conteneurs sans mode privilégié est disponible dans les clusters en version 1.27.4-gke.900 et ultérieure.

Problèmes connus

L'autoscaler de cluster peut mal calculer la capacité des nouveaux nœuds TPU avant que ces nœuds ne signalent les TPU disponibles. L'autoscaler de cluster peut ensuite effectuer un scaling à la hausse supplémentaire et, par conséquent, créer davantage de nœuds que nécessaire. L'autoscaler de cluster effectue un scaling à la baisse des nœuds supplémentaires, s'ils ne sont pas nécessaires, après un fonctionnement à la baisse normal.
L'autoscaler de cluster annule le scaling à la hausse des pools de nœuds TPU qui restent en attente pendant plus de 15 minutes. L'autoscaler de cluster effectue de nouvelles tentatives ultérieurement. Ce comportement peut réduire l'obtention de TPU pour les clients qui n'utilisent pas les réservations.
Les charges de travail non TPU ayant une tolérance pour le rejet TPU peuvent empêcher la réduction de capacité du pool de nœuds si elles sont recréées lors du drainage du pool de nœuds TPU.

Déployer des charges de travail TPU dans GKE

Avant de commencer

Disponibilité des TPU dans GKE

Planifier la configuration TPU

Garantir un quota suffisant

Demander une augmentation du quota de TPU

Créer un cluster

Créer un pool de nœuds

Tranche de TPU à hôte unique

gcloud

Terraform

Console

Tranche de TPU multi-hôte

gcloud

Terraform

Console

Exécuter votre charge de travail sur des nœuds TPU

Préparation des charges de travail

Exemple 1 : Exécuter une charge de travail qui affiche le nombre de puces TPU disponibles dans un pool de nœuds TPU

Exemple 2 : exécuter une charge de travail qui affiche le nombre de puces TPU disponibles dans la VM TPU

Effectuer un nettoyage

Configurations supplémentaires

Multicouche

Migrer votre réservation TPU

Journalisation

Utiliser le provisionnement automatique des nœuds GKE

Réparation automatique des nœuds TPU

Observabilité et métriques

Métriques d'exécution

Métriques d'hôte

Exécuter des conteneurs sans mode privilégié

Problèmes connus

Étapes suivantes