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 Google Cloud é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.
Chaque TensorCore sur une puce TPU est constitué d'une ou de plusieurs unités multiplitives Matrix, d'une unité vectorielle et d'une unité scalaire.
Un MXU est composé de 128 x 128 multiplicateurs/accumulateurs dans un tableau systolique. Les unités UMX fournissent la majeure partie de la puissance de calcul dans un TensorCore. Chaque unité matricielle est capable d'effectuer 16 000 opérations de multiplication dans chaque cycle. Toutes les multiplications utilisent des entrées bfloat16, mais toutes les accumulations sont effectuées au format numérique FP32.
L'unité de vecteur est utilisée pour le calcul général, comme les activations et softmax. 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 du TPU que vous utilisez. Les détails architecturaux et les caractéristiques des performances des TPU v2 et TPU v3 sont disponibles dans la section A Superordinateur propre au domaine pour l'entraînement de réseaux de neurones profonds.
TPU v4
Chaque puce TPU v4 contient deux TensorCores. Chaque TensorCore comporte quatre unités matricielles, une unité vectorielle et une unité scalaire. Le tableau suivant présente les caractéristiques clés d'un pod TPU v4.
| Caractéristiques clés | Valeurs des pods v4 |
|---|---|
| Pic de calcul par chip | 275 téraflops (bf16 ou int8) |
| Capacité et bande passante HBM2 | 32 Gio, 1 200 Gbit/s |
| Puissance minimale/moyenne/maximale | 90/170/192 W |
| Taille du pod TPU | 4096 chips |
| Topologie de l'interconnexion | Tore en 3D |
| Pic de calcul par pod | 1.1 exaflops (bf16 ou int8) |
| Réduire la bande passante par pod | 1,1 Po/s |
| Bande passante bissectionnelle par pod | 24 To/s |
Le schéma suivant illustre une puce TPU v4.
TPU v3
Chaque puce TPU v3 contient deux TensorCores. Chaque TensorCore comporte deux unités MXU, une unité vectorielle et une unité scalaire. Le tableau suivant présente les spécifications clés et leurs valeurs pour un pod TPU v3.
| Caractéristiques clés | Valeurs des pods v3 |
|---|---|
| Pic de calcul par chip | 123 téraflops (bf16) |
| Capacité et bande passante HBM2 | 32 Gio, 900 Gbit/s |
| Puissance minimale/moyenne/maximale | 123/220/262 W |
| Taille du pod TPU | 1024 chips |
| Topologie de l'interconnexion | Torus en 2D |
| Pic de calcul par pod | 126 pétaflops (bf16) |
| Réduire la bande passante par pod | 340 To/s |
| Bande passante bissectionnelle par pod | 6,4 To/s |
Le schéma suivant illustre une puce TPU v3.
TPU v2
La plus petite configuration TPU v2 contient quatre puces TPU et 16 Gio de HBM. Chaque puce TPU contient deux TensorCore. Chaque TensorCore comporte un MXU, une unité vectorielle et une unité scalaire. Le schéma suivant illustre une puce TPU v2.
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 v3
Système de mémoire:
L'architecture de mémoire non uniforme (NUMA) est une architecture de mémoire conçue pour les machines dotées de plusieurs processeurs. Chaque processeur a un accès direct à un bloc de mémoire à grande vitesse. Un CPU et sa mémoire sont appelés un nœud NUMA. NUMA nœuds sont connectés aux NUMA nœuds directement adjacents. Un processeur d'un nœud NUMA peut accéder à la mémoire d'un autre nœud NUMA, mais cet accès est plus lent que l'accès à la mémoire d'un nœud NUMA.
Un logiciel exécuté sur une machine à plusieurs processeurs peut placer les données nécessaires à un processeur au sein de son nœud NUMA, ce qui augmente le débit de la mémoire. Pour en savoir plus sur le NUMA, consultez le document Wikipédia sur l'accès à la mémoire non uniforme.
L'environnement d'exécution TPU v4 n'est pas compatible avec NUMA de manière native. Toutefois, vous pouvez bénéficier des avantages de la NUMA-localité en associant le script d'entraînement au NUMA Node 0.
Pour activer la liaison de nœuds NUMA:
Installez l'outil de ligne de commande numactl.
$ $ sudo apt-get update $ $ sudo apt-get install numactl
Utilisez
numactl --cpunodebind=0lorsque vous lancez votre script d'entraînement. Cette opération associe votre code de script au nœud NUMA 0.$ $ numactl --cpunodebind=0 python3 your-training-script
Activez la liaison de nœuds NUMA dans les cas suivants:
- Si votre charge de travail dépend fortement des charges de travail du processeur (par exemple, classification des images ou charges de travail de recommandation), quel que soit le framework.
- Si vous utilisez une version d'exécution TPU sans suffixe -pod (par exemple, TPU-vm-tf-2.10.0-v4) :
Autres différences entre les systèmes de mémoire:
- Les puces TPU v4 disposent d'un espace mémoire HBM unifié de 32 Gio sur l'ensemble de la puce, ce qui permet une meilleure coordination entre les deux TensorCore sur puces.
- Amélioration des performances de HBM grâce aux normes et vitesses de mémoire les plus récentes.
- Amélioration du profil de performances de la zone de marché désignée avec prise en charge intégrée du mode hautes performances à 512 milliards de précisions.
TensorCore:
- Deux fois plus de unités et une fréquence d'horloge plus élevée, avec 275 TFLOPS max.
- Bande passante de transposition et de permutation x2.
- Modèle d'accès à la mémoire du datastore pour la mémoire partagée (Cmem).
- Bande passante plus rapide pour le chargement MXU et prise en charge du mode 8 bits pour réduire la taille de lot et améliorer la latence d'inférence.
Interconnexion par chip:
Six liens d'interconnexion par puce pour activer des topologies de réseau ayant des diamètres de réseau plus petits.
Autre :
- Interface PCIe x16 PCIe gen3 pour l'hôte (connexion directe).
- Modèle de sécurité amélioré
- Une meilleure efficacité énergétique
Avantages en termes de performances des TPU v3 par rapport aux TPU v2
L'augmentation du nombre de FLOPS par TensorCore et de la capacité de mémoire dans les configurations TPU v3 peut améliorer les performances de vos modèles de différentes manières:
Les configurations TPU v3 offrent d'importants avantages en termes de performances par TensorCore pour les modèles liés au calcul. Les modèles liés à la mémoire sur des configurations TPU v2 peuvent ne pas atteindre cette même amélioration de performances s'ils sont également liés à des configurations TPU v3.
Dans les cas où les données ne rentrent pas dans la mémoire des configurations TPU v2, ce processeur peut 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 des performances du pipeline peut vous aider à résoudre les problèmes de flux d'annonces.
Configurations TPU v4
Un pod TPU v4 est composé de 4 096 puces interconnectées avec des liens reconfigurables à haut débit. La mise en réseau flexible de TPU v4 vous permet de connecter les puces d'une même tranche de plusieurs manières. Lorsque vous créez un secteur TPU, vous spécifiez la version du TPU et le nombre de ressources TPU dont vous avez besoin. Lorsque vous créez un segment TPU v4, vous pouvez spécifier son type et sa taille de deux manières:
Utiliser
AcceleratorTypeCette méthode est recommandée lorsque vous ne spécifiez aucune topologie.
Utiliser
AcceleratorConfigNous vous recommandons d'utiliser
AcceleratorConfiglorsque vous souhaitez personnaliser la topologie physique. Cette étape est généralement requise pour le réglage des performances avec des secteurs >= 256 puces.
Pour configurer des TPU v4 avec AcceleratorType, utilisez l'indicateur --accelerator-type lorsque vous créez votre tranche de TPU. Définissez --accelerator-type sur une chaîne contenant la version de TPU et le nombre de TensorCore que vous souhaitez utiliser. Par exemple, pour créer une tranche v4 avec 32 TensorCore, utilisez --accelerator-type=v4-32.
Pour configurer des TPU v4 avec AcceleratorConfig, utilisez les options --version et --topology. Définissez --version sur la version de TPU que vous souhaitez utiliser et --topology sur la disposition physique des puces TPU de la tranche. Par exemple, pour créer une tranche v4 avec 16 chips, vous devez utiliser --version=v4 et --topology=2x2x4.
La topologie du 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 certaines topologies TPU v4 courantes.
De plus grandes tranches peuvent être créées à partir d'un ou de plusieurs "cubes" de puces 4 x 4.
Les segments TPU d'un nombre donné de puces peuvent être configurés de différentes manières. Par exemple, une tranche de TPU avec 512 puces (1 024 TensorCore) peut être configurée à l'aide des topologies suivantes: 4x4x32, 4x8x16 ou 8x8x8.
Un segment de TPU avec 2 048 puces (4 096 TensorCores) offre encore plus d'options de topologie: 4x4x128, 4x8x64, 4x16x32 et 8x16x16. Pour en savoir plus, consultez la section Types et topologies de TPU.
Configurations des TPU v2 et v3
Les TPU sont disponibles dans les configurations suivantes :
- Carte TPU unique
- Tranche de pod TPU
Carte TPU unique
Une configuration TPU à carte unique est une carte autonome dotée de quatre puces TPU (huit TensorCore) sans connexion réseau à d'autres cartes TPU. Les TPU à carte unique ne font pas partie d'une configuration de pod TPU et n'occupent pas de partie d'un pod TPU.
Tranches de pods TPU
Dans une tranche de pod TPU, les puces TPU sont connectées à l'aide d'une interconnexion haut débit. 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 TensorCore dans une tranche de pod. Les tranches de pod TPU v3 sont disponibles avec les TensorCore 32, 128, 512, 1024 et 2048. Les tranches de pod TPU v2 sont disponibles avec 32, 128, 256 ou 512 TensorCore.
Architectures de VM Cloud TPU
Les TPU sont conçus pour effectuer des opérations matricielles rapidement. Chaque carte TPU est connectée à une machine hôte basée sur un processeur pour effectuer des opérations qui ne peuvent pas être exécutées sur le TPU. Les machines hôtes se chargent de charger les données depuis Cloud Storage, de les prétraiter et de les envoyer au TPU.
Dans un pod TPU, il existe un hôte TPU pour chaque carte TPU.
La manière dont vous interagissez avec l'hôte TPU (et la carte TPU) dépend de l'architecture de la VM TPU que vous utilisez: les nœuds TPU ou les VM TPU.
Nœuds TPU
L'architecture du nœud TPU consiste en 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 est donc difficile de déboguer l'entraînement et les erreurs TPU.
VM TPU
L'architecture de la VM TPU vous permet de vous connecter en SSH à la VM connectée physiquement à l'appareil TPU. Vous disposez d'un accès racine à la VM, ce qui vous permet d'exécuter du code arbitraire. Vous pouvez accéder aux journaux de débogage et aux messages d'erreur du compilateur et de l'environnement d'exécution.
Frameworks compatibles
Des frameworks tels que JAX, PyTorch et TensorFlow accèdent aux TPU à l'aide d'une bibliothèque partagée appelée libtpu, présente sur chaque VM TPU. libtpu inclut le compilateur XLA, le logiciel d'exécution TPU et le pilote TPU.
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.
Pour en savoir plus sur l'exécution de TensorFlow sur Cloud TPU, consultez le guide de démarrage rapide de TensorFlow.
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 sur le réseau gRPC.
PyTorch
VM TPU
Avec les VM TPU, votre code PyTorch s'exécute directement sur les hôtes TPU.
Pour en savoir plus sur PyTorch et Cloud TPU, consultez le guide de l'utilisateur PyTorch/XLA.
Nœuds TPU
PyTorch s'exécute sur l'architecture des nœuds Cloud TPU à l'aide d'une bibliothèque appelée XRT. XRT envoie des graphiques XLA et des instructions d'exécution via gRPC à exécuter sur les serveurs TensorFlow. Une VM utilisateur est requise pour chaque hôte TPU.
JAX
VM TPU
Avec les VM TPU, votre code JAX s'exécute directement sur les hôtes TPU.
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.
