À propos des TPU dans GKE

Cette page présente Cloud TPU avec Google Kubernetes Engine (GKE). Les TPU (Tensor Processing Units) sont des circuits intégrés spécifiques aux applications (ASIC) développés spécifiquement par Google et permettant d'accélérer les charges de travail de machine learning qui utilisent des frameworks tels que TensorFlow, PyTorch et JAX.

Avant d'utiliser des TPU dans GKE, nous vous recommandons de découvrir le fonctionnement des accélérateurs de machine learning avec la présentation de Cloud TPU.

Cette page vous aide à comprendre les principes de base de Cloud TPU avec Google Kubernetes Engine (GKE), y compris la terminologie, les avantages des TPU et les considérations liées à la planification des charges de travail.

Pour apprendre à configurer Cloud TPU dans GKE, consultez les ressources suivantes :

Avantages de l'utilisation des TPU dans GKE

GKE fournit une assistance complète pour la gestion du cycle de vie des nœuds TPU et des pools de nœuds, y compris la création, la configuration et la suppression de VM TPU. GKE est également compatible avec les VM Spot et l'utilisation de Cloud TPU réservés. L'utilisation de TPU dans GKE présente les avantages suivants :

  • Environnement opérationnel cohérent : vous pouvez utiliser une plate-forme unique pour toutes les charges de travail et le machine learning.
  • Mises à jour automatiques : GKE automatise les mises à jour de versions afin de réduire les coûts opérationnels.
  • Équilibrage de charge : GKE répartit la charge afin de réduire la latence et d'améliorer la fiabilité.
  • Scaling réactif : GKE effectue le scaling automatique des ressources TPU pour répondre aux besoins de vos charges de travail.
  • Gestion des ressources : avec Kueue, un système de mise en file d'attente de jobs Kubernetes natif, vous pouvez gérer les ressources de plusieurs locataires au sein de votre organisation à l'aide de la mise en file d'attente, de la préemption, de la priorisation et du partage équitable.

Terminologie relative au TPU dans GKE

Ce document utilise la terminologie suivante liée aux TPU :

  • Type de TPU : type de Cloud TPU, tel que v5e
  • Nœud de tranche de TPU : nœud Kubernetes représenté par une seule VM disposant d'une ou plusieurs puces TPU interconnectées.
  • Pool de nœuds de tranche de TPU : groupe de nœuds Kubernetes au sein d'un cluster qui ont tous la même configuration TPU.
  • Topologie de TPU : nombre et disposition physique des puces TPU dans une tranche de TPU.
  • Atomique : GKE traite tous les nœuds interconnectés comme une seule unité. Lors des opérations de scaling, GKE met à l'échelle l'intégralité de l'ensemble de nœuds à zéro et crée de nouveaux nœuds. Si une machine du groupe échoue ou s'arrête, GKE recrée l'ensemble des nœuds en tant que nouvelle unité.
  • Immuable : vous ne pouvez pas ajouter manuellement de nouveaux nœuds à l'ensemble de nœuds interconnectés. Vous pouvez toutefois créer un pool de nœuds avec la topologie TPU souhaitée et planifier des charges de travail sur ce pool.

Type de pool de nœuds de tranche de TPU

GKE est compatible avec deux types de pools de nœuds TPU :

Le type et la topologie de TPU déterminent si votre nœud de tranche de TPU peut être multi-hôte ou à hôte unique. Notre recommandation :

  • Pour les modèles à grande échelle, utilisez des nœuds de tranche de TPU multi-hôte.
  • Pour les modèles à petite échelle, utilisez des nœuds de tranche de TPU à hôte unique.

Pool de nœuds de tranche de TPU multi-hôte

Un pool de nœuds de tranche de TPU multi-hôte est un pool de nœuds contenant au moins deux VM TPU interconnectées. Chaque VM est associée à un appareil TPU. Les TPU d'une tranche de TPU multi-hôte sont connectés via une interconnexion à haut débit (ICI). Une fois qu'un pool de nœuds de tranche de TPU multi-hôte est créé, vous ne pouvez plus y ajouter de nœuds. Par exemple, vous ne pouvez pas créer un pool de nœuds v4-32, puis ajouter un nœud Kubernetes (VM TPU) supplémentaire au pool de nœuds. Pour ajouter une tranche de TPU supplémentaire à un cluster GKE, vous devez créer un pool de nœuds.

Les VM d'un pool de nœuds de tranche de TPU de multi-hôte sont traitées comme une seule unité atomique. Si GKE ne parvient pas à déployer un nœud dans la tranche, aucun nœud du nœud de tranche de TPU n'est déployé.

Si un nœud d'une tranche de TPU multi-hôte doit être réparé, GKE arrête toutes les VM de la tranche de TPU, ce qui force l'expulsion de tous les pods Kubernetes de la charge de travail. Une fois que toutes les VM de la tranche de TPU sont opérationnelles, les pods Kubernetes peuvent être planifiés sur les VM de la nouvelle tranche de TPU.

Le schéma suivant illustre une tranche de TPU multi-hôte v5litepod-16 (v5e). Cette tranche de TPU comporte quatre VM. Chaque VM de la tranche de TPU comporte quatre puces TPU v5e connectées par des interconnexions à haut débit (ICI), et chaque puce TPU v5e comporte un TensorCore.

Schéma de tranche de TPU multi-hôte

Le schéma suivant montre un cluster GKE contenant une tranche de TPU v5litepod-16 (v5e, topologie : 4x4) et une tranche de TPU v5litepod-8 (v5e, topologie : 2x4) :

Schéma de pod TPU v5e

Pools de nœuds de tranche de TPU à hôte unique

Un pool de nœuds de tranche à hôte unique est un pool de nœuds contenant une ou plusieurs VM TPU indépendantes. Chaque VM est associée à un appareil TPU. Bien que les VM d'un pool de nœuds de tranche à hôte unique puissent communiquer via le réseau de centre de données (DCN), les TPU associés aux VM ne sont pas interconnectés.

Le schéma suivant montre un exemple de tranche de TPU à hôte unique contenant sept machines v4-8 :

Schéma du pool de nœuds de tranche à hôte unique

Caractéristiques des TPU dans GKE

Les TPU présentent des caractéristiques uniques qui nécessitent une planification et une configuration spécifiques.

Topologie

La topologie définit la disposition physique des TPU dans une tranche de pod TPU. GKE provisionne une tranche de pod TPU dans des topologies à deux ou trois dimensions en fonction de la version du TPU. Vous spécifiez une topologie en tant que nombre de puces TPU dans chaque dimension, comme suit :

Pour un TPU v4 et v5p planifié dans des pools de nœuds TPU à hôtes multiples, vous définissez la topologie à trois tuples ({A}x{B}x{C}), par exemple 4x4x4. Le produit de {A}x{B}x{C} définit le nombre de puces TPU dans le pool de nœuds. Par exemple, vous pouvez définir des topologies de petite taille qui comportent moins de 64 puces TPU avec des formes de topologie telles que 2x2x2, 2x2x4 ou 2x4x4. Si vous utilisez des topologies de plus de 64 puces TPU, les valeurs que vous attribuez à {A}, {B} et {C} doivent répondre aux conditions suivantes :

  • {A}, {B} et {C} doivent être des multiples de quatre.
  • La topologie la plus grande compatible avec v4 est 12x16x16 et avec v5p 16x16x24.
  • Les valeurs attribuées doivent conserver le modèle A ≤ B ≤ C. Par exemple, 4x4x8 ou 8x8x8.

Type de machine

Les types de machines compatibles avec les ressources TPU suivent une convention d'attribution de noms qui inclut la version de TPU et le nombre de puces par nœud, par exemple ct<version>-hightpu-<node-chip-count>t. Par exemple, le type de machine ct5lp-hightpu-1t est compatible avec les TPU v5e et ne contient qu'une seule puce TPU.

Mode privilégié

Le mode privilégié remplace de nombreux autres paramètres de sécurité dans securityContext. Pour accéder aux TPU, les conteneurs s'exécutant dans des nœuds GKE dans :

  • Les versions 1.28 et antérieures doivent activer le mode privilégié.
  • Les versions 1.28 ou ultérieures n'ont pas besoin du mode privilégié.

Fonctionnement des TPU dans GKE

La gestion et la hiérarchisation des ressources Kubernetes traitent les VM sur les TPU de la même manière que les autres types de VM. Pour demander des puces TPU, utilisez le nom de ressource google.com/tpu :

    resources:
        requests:
          google.com/tpu: 4
        limits:
          google.com/tpu: 4

Lorsque vous utilisez des TPU dans GKE, vous devez prendre en compte les caractéristiques TPU suivantes :

  • Une VM peut accéder à huit puces TPU au maximum.
  • Une tranche de pod TPU contient un nombre fixe de puces TPU, qui dépend du type de machine TPU que vous choisissez.
  • Le nombre de google.com/tpu demandés doit être égal au nombre total de puces TPU disponibles sur le nœud de tranche de TPU. Tout conteneur d'un pod GKE qui demande des TPU doit consommer toutes les puces TPU du nœud. Sinon, votre déploiement échoue, car GKE ne peut pas consommer partiellement les ressources TPU. Étudions les cas de figure suivants :
    • Le type de machine ct5l-hightpu-8t possède un seul nœud de tranche TPU avec huit puces TPU. Par conséquent, sur un nœud :
      • Vous pouvez déployer un pod GKE nécessitant huit puces TPU.
      • Vous ne pouvez pas déployer deux pods GKE nécessitant chacun quatre puces TPU.
    • Le type de machine ct5lp-hightpu-4t avec une topologie 2x4 contient deux nœuds de tranche de TPU avec quatre puces TPU chacun, soit un total de huit puces TPU. Avec ce type de machine :
      • Vous ne pouvez pas déployer un pod GKE nécessitant huit puces TPU sur les nœuds de ce pool de nœuds.
      • Vous pouvez déployer deux pods nécessitant chacun quatre puces TPU, avec un pod sur chaque nœud de ce pool de nœuds.
    • Le TPU v5e avec la topologie 4x4 comporte 16 puces TPU dans quatre nœuds. La charge de travail GKE Autopilot qui sélectionne cette configuration doit demander quatre puces TPU dans chaque instance dupliquée, pour une à quatre instances répliquées.
  • Dans les clusters Standard, plusieurs pods Kubernetes peuvent être programmés sur une VM, mais un seul conteneur de chaque pod peut accéder aux puces TPU.
  • Pour créer des pods kube-system, tels que kube-dns, chaque cluster Standard doit disposer d'au moins un pool de nœuds de tranche non TPU.
  • Par défaut, les nœuds de tranche de TPU ont un rejet google.com/tpu qui empêche la planification des charges de travail non TPU sur les nœuds de tranche de TPU. Les charges de travail qui n'utilisent pas de TPU sont exécutées sur des nœuds non TPU, ce qui libère de la capacité de calcul sur les nœuds de tranche de TPU pour le code qui utilise des TPU. Notez que le rejet ne garantit pas que les ressources TPU sont entièrement utilisées.
  • GKE collecte les journaux émis par les conteneurs s'exécutant sur des nœuds de tranche de TPU. Pour en savoir plus, consultez la section Journalisation.
  • Les métriques d'utilisation du TPU, telles que les performances d'exécution, sont disponibles dans Cloud Monitoring. Pour en savoir plus, consultez la section Observabilité et métriques.

Étape suivante