Architecture du système

Les TPU (Tensor Processing Units) sont des accélérateurs de ML conçus par Google. Cloud TPU rend les TPU disponibles en tant que ressource cloud GCP évolutive. Vous pouvez exécuter des charges de travail de machine learning sur des Cloud TPU à l'aide de frameworks de machine learning tels que TensorFlow, Pytorch et JAX.

L'unité Matrix-multiply est composée de 128 x 128 multiplicateurs et accumulateurs dans un tableau systolique. Les MXU fournissent l'essentiel de la puissance de calcul dans un TensorCore (TPU Core). Chaque unité matricielle est capable d'effectuer 16 000 opérations multiply-accumulate par cycle en utilisant le format numérique bfloat16.

L'unité vectorielle est utilisée pour les calculs généraux tels que les activations, softmax, etc. L'unité scalaire est utilisée pour le flux de contrôle, le calcul d'adresses mémoire et d'autres opérations de maintenance.

Versions de TPU

La mise en page exacte d'un TPU dépend de la version de TPU que vous utilisez. Les détails de l'architecture et les caractéristiques des performances de TPU v2 et v3 sont disponibles dans un superordinateur spécifique au domaine pour l'entraînement de réseaux de neurones profonds.

TPU v4

La plus petite configuration TPU v4 contient quatre puces TPU, 32 Gio de mémoire HBM et 128 Mio de mémoire commune partagée. Chaque puce TPU contient deux TensorCores. Chaque TensorCore comporte quatre unités MXU, une unité vectorielle et une unité scalaire. Le diagramme suivant illustre une puce TPU v4.

image

TPU v3

La plus petite configuration TPU v3 contient quatre puces TPU et 32 Gio de mémoire HBM. Chaque puce TPU contient deux TensorCores. Chaque TensorCore comporte deux unités MXU, une unité vectorielle et une unité scalaire. Le schéma suivant illustre une puce TPU v3.

image

TPU v2

La plus petite configuration TPU v2 contient quatre puces TPU et 16 Gio de mémoire HBM. Chaque puce TPU contient deux TensorCores. Chaque TensorCore comporte un MXU, une unité vectorielle et une unité scalaire. Le schéma suivant illustre une puce TPU v2.

image

Cloud TPU propose les configurations de TPU suivantes :

  • TPU individuel
  • Pod TPU : groupe d'appareils TPU connectés par des interconnexions à haut débit
  • Tranche de TPU : subdivision d'un pod TPU

Avantages de TPU v4 par rapport à la version 3

Système de mémoire:

  • Les puces TPU v4 disposent d'un espace de mémoire HBM unifié de 32 Gio sur l'ensemble de la puce, ce qui permet une meilleure coordination entre les deux puces TensorCore.
  • Amélioration des performances de HBM grâce aux dernières vitesses et normes de mémoire.
  • Amélioration du profil de performances de la zone de marché désignée avec compatibilité native pour l'association hautes performances avec 512 milliards de précisions.

TensorCores:

  • Doubler le nombre d'unités d'unité continue et un taux d'horloge plus élevé, fournissant jusqu'à 275 TFLOPS max.
  • Bande passante et permutation en deux fois.
  • Modèle d'accès à la mémoire en magasin (Cmem).
  • Bande passante de chargement MXU plus rapide et mode 8 bits compatibles avec des tailles de lot inférieures et une latence d'inférence améliorée.

Interconnexion puce:

Six liens d'interconnexion par puce pour activer les topologies de réseaux ayant des diamètres plus petits.

Autre :

  • Interface PCIe x16 gen3 gen à l'hôte (connexion directe).
  • Modèle de sécurité amélioré.
  • Plus d'efficacité énergétique

Avantages en termes de performances des TPU v3 par rapport aux TPU v2

L'augmentation de FLOPS par cœur et de la capacité de mémoire accrue dans les configurations TPU v3 peut améliorer les performances de vos modèles des manières suivantes :

  • Les configurations TPU v3 offrent des avantages significatifs en termes de performances par cœur pour les modèles subordonnés aux calculs. Les modèles subordonnés à la mémoire sur les configurations TPU v2 peuvent ne pas bénéficier de cette même amélioration de performances s'ils sont également subordonnés à la mémoire sur les configurations TPU v3.

  • En cas de capacité de mémoire insuffisante sur les configurations TPU v2, les TPU v3 peuvent améliorer les performances et réduire le recalcul des valeurs intermédiaires (rematérialisation).

  • Les configurations TPU v3 peuvent exécuter de nouveaux modèles avec des tailles de lots supérieures à celles acceptées par les configurations TPU v2. Par exemple, les TPU v3 acceptent des réseaux ResNet plus profonds et des images plus grandes avec RetinaNet.

Les modèles qui, avec un TPU v2, sont presque subordonnés au flux d'entrée ("infeed") parce que certaines étapes d'entraînement attendent une entrée, peuvent également être subordonnés aux entrées avec Cloud TPU v3. Le guide pour un pipeline performant peut vous aider à résoudre les problèmes de flux d'entrée.

Configurations TPU v4

Un pod TPU v4 est composé de 4 096 puces interconnectées à interconnexions haut débit reconfigurables. La flexibilité du réseau TPU v4 et #39 vous permet de connecter les puces dans une tranche de même taille de plusieurs manières. Vous allez configurer des TPU v4 à l'aide d'AcceleratorConfig. L'outil AcceleratorConfig vous permet de spécifier la taille des tranches en termes de chips dans le champ de topologie.

La topologie TPU v4 est spécifiée à l'aide d'un 3 tuple qui décrit le nombre, la forme et les interconnexions entre les puces TPU. Les illustrations suivantes présentent des topologies courantes de TPU v4.

image

Les tranches de plus grande taille peuvent être créées à partir d'un ou de plusieurs cubes de 4 x 4 pixels.

Les tranches de TPU d'un nombre donné de puces peuvent être configurées de différentes manières. Par exemple, une tranche de TPU avec le type Accelerator v4-1024 peut désormais être configurée comme suit: 4x4x32, 4x8x16, 8x8x8. Un TPU de la version v4-4096 offre encore plus de flexibilité : 4x4 x 128, 4 x 8 x 64, 4 x 16 x 32, 8 x 16 x 16, 8 x 16 x 16_twistés. Pour en savoir plus, consultez la page Types et topologies TPU.

Configurations TPU v2 et v3

Les TPU sont disponibles dans les configurations suivantes :

  1. Carte TPU unique
  2. Pod TPU
  3. Tranche de pod TPU

Carte TPU unique

Une configuration TPU à carte unique est une carte autonome avec quatre puces TPU (8 TensorCores) sans connexion réseau à d'autres cartes TPU. Les cartes TPU uniques ne font pas partie d'une configuration de pod TPU et ne constituent pas une partie d'un pod TPU.

Pods TPU et tranches de pods TPU

Dans un pod TPU ou une tranche de pod TPU, les puces TPU sont connectées via une interconnexion à haut débit, et chaque puce TPU communique directement avec les autres puces de l'appareil TPU. Le logiciel TPU gère automatiquement la distribution des données à chaque cœur TPU dans un pod ou une tranche. Les tranches de pod sont disponibles avec 32, 128, 512, 1 024 ou 2 048 cœurs.

Architectures de VM Cloud TPU

Les TPU ne peuvent effectuer que des opérations matricielles. Par conséquent, pour effectuer des opérations qui ne peuvent pas être exécutées sur le TPU, chaque carte TPU est connectée à une machine hôte basée sur un processeur. Les machines hôtes sont responsables du chargement des données depuis Cloud Storage, du prétraitement des données et de l'envoi de données au TPU.

image

Dans un pod TPU, il existe un hôte TPU pour chaque carte TPU.

image

La manière dont vous interagissez avec l'hôte TPU (et la carte TPU) dépend de l'architecture de VM TPU que vous utilisez : nœuds TPU ou VM TPU.

Nœuds TPU

L'architecture de nœud TPU se compose d'une VM utilisateur qui communique avec l'hôte TPU via gRPC. Lorsque vous utilisez cette architecture, vous ne pouvez pas accéder directement à l'hôte TPU. Il peut être difficile de déboguer les erreurs d'entraînement et de TPU.

image

VM TPU

L'architecture de la VM TPU supprime le besoin de la VM utilisateur et vous pouvez vous connecter directement à la VM connectée à l'appareil TPU. Vous avez un accès racine à la VM. Vous pouvez donc exécuter du code arbitraire. Vous pouvez accéder aux journaux et aux messages d'erreur du compilateur et d'exécution.

image

Frameworks compatibles

Les frameworks tels que JAX, PyTorch et TensorFlow accèdent aux TPU via une bibliothèque partagée appelée libtpu, qui est présente sur chaque VM TPU. Cette bibliothèque inclut le compilateur XLA utilisé pour compiler les programmes TPU, le logiciel TPU utilisé pour exécuter les programmes compilés et le pilote TPU utilisé par l'environnement d'exécution pour un accès de bas niveau au TPU.

image

TensorFlow

VM TPU

Avec les VM TPU, au lieu d'exécuter votre code Python sur une VM utilisateur, celui-ci peut s'exécuter directement sur l'hôte TPU.

image

Pour en savoir plus sur TensorFlow et Cloud TPU, consultez la page Exécuter des modèles TensorFlow sur Cloud TPU.

Nœuds TPU

À l'origine, l'architecture système des nœuds Cloud TPU a été conçue pour TensorFlow. Les hôtes TPU sont inaccessibles à l'utilisateur et exécutent une copie sans adresse IP du serveur TensorFlow. Ils n'exécutent aucun code Python ni aucun code utilisateur qui n'est pas représenté sous forme de graphe TensorFlow. Le code utilisateur s'exécute dans une VM utilisateur distincte qui communique avec les hôtes TPU via le réseau gRPC.

image

PyTorch

VM TPU

Avec les VM TPU, votre code PyTorch s'exécute directement sur les hôtes TPU.

images

Pour en savoir plus sur PyTorch et Cloud TPU, consultez la page Exécuter des modèles PyTorch sur Cloud TPU.

Nœuds TPU

PyTorch s'exécute sur l'architecture de nœuds Cloud TPU à l'aide d'une bibliothèque appelée XRT, qui permet d'envoyer des graphes XLA et des instructions d'exécution via des connexions gRPC TensorFlow et de les exécuter sur les serveurs TensorFlow. Une VM utilisateur est requise pour chaque hôte TPU.

image

JAX

VM TPU

Avec les VM TPU, les VM utilisateur ne sont pas nécessaires, car vous pouvez exécuter votre code directement sur les hôtes TPU.

image

Pour en savoir plus sur l'exécution de JAX sur Cloud TPU, consultez le guide de démarrage rapide de JAX.

Nœuds TPU

Sur les nœuds Cloud TPU, JAX s'exécute de la même manière que PyTorch, dans la mesure où une VM utilisateur distincte est requise pour chaque VM hôte.

image

Étapes suivantes