Cette page explique comment demander et déployer des charges de travail d'entraînement, de réglage et d'inférence à grande échelle pour les modèles d'intelligence artificielle (IA) et de machine learning (ML) à l'aide d'accélérateurs Cloud TPU (TPU) dans des clusters Google Kubernetes Engine (GKE) Standard.
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
.
Disponibilité des TPU dans GKE
Utilisez GKE pour créer et gérer des pools de nœuds avec des 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 dont il a besoin. Voici les étapes pertinentes lors de la planification de votre configuration TPU :
Garantir un quota suffisant pour les VM Spot ou à la demande
Si vous créez un pool de nœuds de tranche TPU avec des VM à la demande ou Spot, vous devez disposer d'un quota de TPU suffisant dans la région que vous souhaitez utiliser.La création d'un pool de nœuds de tranche TPU qui consomme une réservation TPU ne nécessite pas de quota de TPU1. Vous pouvez ignorer cette étape en toute sécurité pour les TPU réservés.
La création d'un pool de nœuds de tranche TPU à la demande ou Spot dans GKE nécessite un quota d'API Compute Engine. Le quota de l'API Compute Engine (compute.googleapis.com) est différent de celui de l'API Cloud TPU (tpu.googleapis.com), qui est nécessaire lors de la création de TPU avec l'API Cloud TPU.
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.
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é Type et choisissez Quota.
Sélectionnez la propriété Nom, 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-
, saisissezTPU 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 Spot2 TPU v4 ct4p-
TPU v4 PodSlice chips
Preemptible TPU v4 PodSlice chips
TPU v5e ct5l-
TPU v5 Lite Device chips
Preemptible TPU v5 Lite Device chips
TPU v5e ct5lp-
TPU v5 Lite PodSlice chips
Preemptible TPU v5 Lite PodSlice chips
TPU v5p ct5p-
TPU v5p chips
Preemptible TPU v5p chips
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, saisissezregion:us-west4
si vous envisagez de créer des nœuds de tranche TPU dans la zoneus-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.
-
Lorsque vous créez un pool de nœuds de tranche TPU, utilisez les options
--reservation
et--reservation-affinity=specific
pour créer une instance réservée. Les réservations TPU sont disponibles lors de la souscription d'un engagement. ↩ -
Lorsque vous créez un pool de nœuds de tranche TPU, utilisez l'option
--spot
pour créer une instance Spot. ↩
S'assurer de la disponibilité des réservations
La création d'un pool de nœuds de tranche TPU réservé, qui consomme une réservation, ne nécessite aucun quota TPU. Toutefois, la réservation doit comporter suffisamment de puces TPU disponibles ou non utilisées au moment de la création du pool de nœuds.
Pour afficher les réservations d'un projet, affichez la liste de vos réservations.
Pour connaître le nombre de puces TPU disponibles dans une réservation TPU, affichez les détails d'une réservation.
Créer un cluster
Créez un cluster GKE en mode Standard dans une région où des TPU sont 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 \
--cluster-version VERSION
Remplacez les éléments suivants :
CLUSTER_NAME
: nom du nouveau clusterLOCATION
: région dans laquelle votre capacité de TPU est disponible.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 connaître les versions minimales de GKE disponibles par type de machine TPU, consultez la page 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
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-
, utilisezus-central1-a
oueurope-west4-b
. - Pour les types de machines TPU v5e commençant par
ct5lp-
, utilisezus-west1-c
,us-west4-a
,us-west4-b
,us-central1-a
,us-east1-c
,us-east5-b
oueurope-west4-a
. - Pour les TPU v5p, utilisez
us-east1-d
,us-east5-a
ouus-east5-c
.
Pour en savoir plus, consultez la section Disponibilité des TPU dans GKE.
- Pour les TPU v4, utilisez
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=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 est3
. Si l'autoscaling est activé pour le pool de nœuds à l'aide de l'option--enable-autoscaling
, nous vous recommandons de définirNUM_NODES
sur0
, car l'autoscaler provisionne des nœuds supplémentaires dès que vos charges de travail en ont besoin.--reservation=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 section Réservations TPU.--enable-autoscaling
: Ajouter un pool de nœuds avec l'autoscaling activé. Requiert les options supplémentaires suivantes :--total-min-nodes=TOTAL_MIN_NODES
: nombre minimal de nœuds dans le pool de nœuds.--total-max-nodes=TOTAL_MAX_NODES
: nombre maximal de nœuds dans le pool de nœuds.--location-policy=ANY
: donne la priorité à l'utilisation des réservations inutilisées et réduit le risque de préemption des VM Spot.
--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.
Pour obtenir la liste complète de toutes les options que vous pouvez spécifier, consultez la documentation de référence sur gcloud container clusters create
.
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]
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 projetCLUSTER_NAME
: nom du cluster existant.POOL_NAME
: nom du pool de nœuds à créerCLUSTER_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 section 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.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 une tranche de TPU à hôte unique, GKE effectue l'ajustement entre les valeursTOTAL_MIN_NODES
etTOTAL_MAX_NODES
.TOTAL_MIN_NODES
: nombre minimal de nœuds dans le pool de nœuds. Ce champ est facultatif, sauf si l'auto-scaling 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'auto-scaling 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 deRESERVATION_LABEL_VALUES
dans le champreservation_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.
Dans la liste des clusters, cliquez sur le nom du cluster que vous souhaitez modifier.
Cliquez sur add_box 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-
, utilisezus-central1-a
oueurope-west4-b
. - Pour les types de machines TPU v5e commençant par
ct5lp-
, utilisezus-west1-c
,us-west4-a
,us-west4-b
,us-central1-a
,us-east1-c
,us-east5-b
oueurope-west4-a
. - Pour les TPU v5p, utilisez
us-east1-d
,us-east5-a
ouus-east5-c
.
- Pour les TPU v4, utilisez
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
- CT5P : TPU v5p
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 de tranche 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 afin d'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
. - Les types de machines TPU v5e commençant par
ct5l-
ne sont jamais à hôtes multiples. - Pour les types de machines TPU v5e commençant par
ct5lp-
, utilisezus-west1-c
,us-west4-a
,us-west4-b
,us-central1-a
,us-east1-c
,us-east5-b
oueurope-west4-a
. - Pour les types de machines TPU v5p commençant par
ct5p-
, utilisezus-east1-d
,us-east5-a
ouus-east5-c
.
Pour en savoir plus, consultez la section Disponibilité des TPU dans GKE.
- Pour les TPU v4, utilisez
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 :- TPU v4 ou v5p : définissez la topologie à trois tuples (
{A}x{B}x{C}
), par exemple4x4x4
. - TPU v5e : définissez la topologie à deux tuples (
{A}x{B}
), par exemple2x2
.
Pour en savoir plus, consultez la section Topologie.
- TPU v4 ou v5p : définissez la topologie à trois tuples (
NUM_NODES
: Nombre de nœuds dans le pool de nœuds. Doit être égal à zéro ou au produit des valeurs définies dansTPU_TOPOLOGY
({A}x{B}x{C}
) divisé 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 votreTPU_TOPOLOGY
est2x4x4
(TPU v4 avec quatre puces dans chaque VM), la valeur deNUM_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 de tranche 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 de tranche 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. L'option--max-nodes
est obligatoire si--enable-autoscaling
est fourni. Elle doit être égale au produit des valeurs définies dansTPU_TOPOLOGY
({A}x{B}x{C}
) divisée par le nombre de puces dans pour 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 projetCLUSTER_NAME
: nom du cluster existant auquel ajouter le pool de nœuds.POOL_NAME
: nom du pool de nœuds à créerCLUSTER_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 section 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 de tranche TPU comporte quatre puces. Par exemple, siTPU_TOPOLOGY
est défini sur4x8
, il y a 32 puces, ce qui signifie queNUM_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 exemple4x4x4
. - TPU v5e : définissez la topologie à deux tuples (
{A}x{B}
), par exemple2x2
.
- Pour les TPU v4 : définissez la topologie à trois tuples (
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 deRESERVATION_LABEL_VALUES
dans le champreservation_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 dansTPU_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 de tranche 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.
Dans la liste des clusters, cliquez sur le nom du cluster que vous souhaitez modifier.
Cliquez sur add_box 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
. - Les types de machines TPU v5e commençant par
ct5l-
ne sont jamais à hôtes multiples. - Pour les types de machines TPU v5e commençant par
ct5lp-
, utilisezus-west1-c
,us-west4-a
,us-west4-b
,us-central1-a
,us-east1-c
,us-east5-b
oueurope-west4-a
. - Pour les types de machines TPU v5p commençant par
ct5p-
, utilisezus-east1-d
,us-east5-a
ouus-east5-c
.
- Pour les TPU v4, utilisez
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 de tranche 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.
État de provisionnement
Si GKE ne peut pas créer votre pool de nœuds de tranche TPU en raison d'une capacité TPU insuffisante, il renvoie un message d'erreur indiquant que les nœuds de tranche TPU ne peuvent pas être créés en raison d'un manque de capacité.
Si vous créez un pool de nœuds de tranche TPU à hôte unique, le message d'erreur se présente comme suit :
2 nodes cannot be created due to lack of capacity. The missing nodes will be
created asynchronously once capacity is available. You can either wait for the
nodes to be up, or delete the node pool and try re-creating it again later.
Si vous créez un pool de nœuds de tranche TPU multi-hôte, le message d'erreur se présente comme suit :
The nodes (managed by ...) cannot be created now due to lack of capacity. They
will be created asynchronously once capacity is available. You can either wait
for the nodes to be up, or delete the node pool and try re-creating it again
later.
Votre requête de provisionnement TPU peut rester dans la file d'attente pendant une longue période et reste à l'état "Provisionnement" tant qu'elle est dans la file d'attente.
Une fois la capacité disponible, GKE crée les nœuds restants qui n'ont pas été créés.
Si vous avez besoin de capacité plus tôt, envisagez d'essayer des VM Spot. Notez toutefois que les VM Spot consomment un quota différent de celui des instances à la demande.
Vous pouvez supprimer la requête TPU en file d'attente en supprimant le pool de nœuds du tranche TPU.
Exécuter votre charge de travail sur des nœuds de tranche TPU
Préparation de la charge de travail
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 delibtpu.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ètreTPU_WORKER_ID
.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-
, utiliseztpu-v5-lite-device
. - Pour les types de machines TPU v5e commençant par
ct5lp-
, utiliseztpu-v5-lite-podslice
. - Pour les TPU v5p, utilisez
tpu-v5p-slice
.
- Pour les TPU v4, utilisez
TPU_TOPOLOGY
: topologie physique de la tranche 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 exemple4x4x4
. - TPU v5e : définissez la topologie à deux tuples (
{A}x{B}
), par exemple2x2
. - TPU v5p : définissez la topologie à trois tuples (
{A}x{B}x{C}
), par exemple4x4x4
.
- TPU v4 : définissez la topologie à trois tuples (
GKE injecte automatiquement ces variables d'environnement en utilisant un webhook en mutation lorsqu'un job est créé avec
Lors du déploiement de ressources TPU multi-hôtes avec Kuberay, GKE fournit un webhook déployable dans le cadre des modèles Terraform pour exécuter Ray sur GKE. Les instructions d'exécution de Ray sur GKE avec des TPU sont disponibles dans le guide de l'utilisateur TPU. Le webhook à mutation injecte ces variables d'environnement dans les clusters Ray qui demandent des propriétéscompletionMode: Indexed
,subdomain
,parallelism > 1
et en demandant des propriétésgoogle.com/tpu
. GKE ajoute un service sans tête afin que les enregistrements DNS soient ajoutés pour les pods qui sauvegardent le service.google.com/tpu
et un sélecteur de nœudcloud.google.com/gke-tpu-topology
à plusieurs hôtes.
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 de tranche TPU
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. 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: 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 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 tranche TPU
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. 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
- 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
Mettre à niveau des pools de nœuds à l'aide d'accélérateurs (GPU et TPU)
GKE met automatiquement à niveau les clusters standards, y compris les pools de nœuds. Vous pouvez également mettre à jour manuellement les pools de nœuds si vous souhaitez que vos nœuds passent plus rapidement à la nouvelle version. Pour contrôler le fonctionnement des mises à niveau de votre cluster, utilisez les versions disponibles, les intervalles et exclusions de maintenance et le séquençage du déploiement.
Vous pouvez également configurer une stratégie de mise à niveau des nœuds pour votre pool de nœuds, telle que des mises à niveau de la surutilisation ou des mises à niveau bleu-vert. En configurant ces stratégies, vous pouvez vous assurer que les pools de nœuds sont mis à niveau de manière à trouver l'équilibre optimal entre la vitesse et les interruptions pour votre environnement. Pour les pools de nœuds de tranches TPU multi-hôtes, au lieu d'utiliser la stratégie de mise à niveau des nœuds configurée, GKE recrée de manière atomique l'ensemble du pool de nœuds en une seule étape. Pour en savoir plus, consultez la définition de l'atomicité dans la section Terminologie relative au TPU dans GKE.
L'utilisation d'une stratégie de mise à niveau des nœuds nécessite temporairement que GKE provisionne des ressources supplémentaires, en fonction de la configuration. Si Google Cloud dispose d'une capacité limitée pour les ressources de votre pool de nœuds, par exemple, si des erreurs de disponibilité des ressources s'affichent lorsque vous essayez de créer d'autres nœuds avec des GPU ou des TPU, consultez la section Effectuer une mise à niveau dans un environnement à ressources limitées.
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 de tranche TPU qui n'ont plus de charges de travail programmées. Si les charges de travail en cours d'exécution doivent être arrêtées progressivement, utilisez kubectl drain
pour nettoyer les charges de travail avant de supprimer le nœud.
Supprimez un pool de nœuds de tranche 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 avez déjà effectué des réservations TPU, vous devez d'abord les migrer 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 sans avoir à effectuer de migration. Pour savoir comment migrer vos réservations TPU, consultez la section 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 de tranche TPU
Si un nœud de tranche TPU dans un pool de nœuds de tranche TPU multi-hôte n'est pas opérationnel, l'intégralité du pool de nœuds est recréé. Dans un pool de nœuds de tranche TPU à hôte unique, seul le nœud TPU non opérationnel est réparé automatiquement.
Les conditions qui entraînent des nœuds de tranche TPU non opérationnels sont les suivantes :
- Tout nœud de tranche TPU avec des conditions de nœud communes.
- Tout nœud de tranche TPU dont le nombre de TPU non allouables est supérieur à zéro.
- Toute instance de VM dans une tranche TPU arrêtée (en raison de la préemption) ou terminée.
- Maintenance des nœuds : si un nœud de tranche TPU d'un pool de nœuds de tranche TPU multi-hôte est défaillant pour la maintenance de l'hôte, GKE recrée l'ensemble du pool de nœuds de tranche TPU.
Vous pouvez consulter l'état de la réparation (y compris la raison de l'échec) dans l'historique des opérations. Si l'échec est causé par un quota insuffisant, contactez votre responsable de compte Google Cloud pour augmenter le quota correspondant.
Configurer l'arrêt optimal des nœuds de tranche TPU
Dans les clusters GKE dont le plan de contrôle exécute la version 1.29.1-gke.1425000 ou ultérieure, les nœuds de tranche TPU acceptent les signaux SIGTERM
qui avertissent le nœud d'un arrêt imminent. La notification d'arrêt imminent peut être configurée jusqu'à cinq minutes dans les nœuds TPU.
Pour configurer GKE afin qu'il arrête correctement vos charges de travail dans ce délai de notification, suivez la procédure décrite dans la section Gérer les interruptions des nœuds GKE pour les GPU et les TPU.
Exécuter des conteneurs sans mode privilégié
Si le nœud de tranche TPU exécute une version antérieure à 1.28, lisez la section suivante :
Un conteneur s'exécutant sur une VM dans une tranche de 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 un ulimit
plus élevé. 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.
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.
Problèmes connus
- L'autoscaler de cluster peut mal calculer la capacité des nouveaux nœuds de tranche TPU avant que ces nœuds ne signalent les TPU disponibles. L'autoscaler de cluster peut alors effectuer un scaling à la hausse supplémentaire et créer ainsi plus de nœuds que nécessaire. L'autoscaler de cluster effectue le scaling à la baisse des nœuds supplémentaires, s'ils ne sont pas nécessaires, après une opération de scaling à la baisse standard.
- L'autoscaler de cluster annule le scaling à la hausse des pools de nœuds de tranche TPU qui restent en attente pendant plus de 10 heures. L'autoscaler de cluster relance ces opérations de scaling à la hausse ultérieurement. Ce comportement peut réduire la capacité d'obtention des 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 de tranche TPU.
- La métrique d'utilisation de la bande passante de mémoire n'est pas disponible pour les TPU v5e.
Étape suivante
- Diffuser des grands modèles de langage avec Saxml sur des TPU
- En savoir plus sur la configuration de Ray sur GKE avec des TPU
- Créer des modèles de machine learning à grande échelle dans Cloud TPU avec GKE
- Diffuser des grands modèles de langage avec KubeRay sur des TPU
- Résoudre les problèmes liés aux TPU dans GKE