Architecture des TPU
Les TPU (Tensor Processing Units) sont des circuits intégrés propres aux applications (ASIC) conçus par Google pour accélérer les charges de travail de machine learning. Cloud TPU est un Google Cloud service qui met à disposition les TPU en tant que ressource évolutive.
Les TPU sont conçus pour effectuer rapidement des opérations matricielles, ce qui les rend idéaux pour les charges de travail de machine learning. Vous pouvez exécuter des charges de travail de machine learning sur des TPU à l'aide de frameworks tels que Pytorch et JAX.
Comment fonctionnent les TPU ?
Pour comprendre le fonctionnement des TPU, il est utile de comprendre comment d'autres accélérateurs gèrent les difficultés de calcul liées à l'entraînement des modèles de ML.
Fonctionnement d'un processeur
Un processeur est un processeur à usage général basé sur l'architecture de von Neumann. Cela signifie qu'un processeur fonctionne avec un logiciel et une mémoire, comme illustré ci-dessous :
Le principal avantage des processeurs est leur flexibilité. Vous pouvez charger n'importe quel type de logiciel sur un processeur pour de nombreux types d'applications. Par exemple, vous pouvez utiliser un processeur pour le traitement de texte sur un PC, contrôler les moteurs de fusées, exécuter des transactions bancaires ou classer des images avec un réseau de neurones.
Un processeur charge les valeurs à partir de la mémoire, effectue un calcul sur les valeurs et stocke le résultat en mémoire pour chaque calcul. L'accès à la mémoire est lent par rapport à la vitesse de calcul et peut limiter le débit total des processeurs. On parle souvent de goulot d'étranglement de von Neumann.
Fonctionnement d'un GPU
Pour obtenir un débit plus élevé, les GPU contiennent des milliers d'unités logiques arithmétiques (ALU) dans un seul processeur. Un GPU moderne contient généralement entre 2 500 et 5 000 ALU. Grâce à ce grand nombre de processeurs, vous pouvez exécuter des milliers de multiplications et d'additions simultanément.
Cette architecture GPU fonctionne bien sur les applications à parallélisme massif, comme les opérations matricielles dans un réseau de neurones. En réalité, avec une charge de travail d'entraînement classique pour le deep learning, un GPU peut fournir une commande beaucoup plus importante qu'un processeur.
Toutefois, le GPU reste un processeur à usage général qui doit prendre en charge de nombreuses applications et logiciels différents. Par conséquent, les GPU rencontrent le même problème que les processeurs. Pour chaque calcul dans les milliers d'ALU, un GPU doit accéder aux registres ou à la mémoire partagée pour lire les opérandes et stocker les résultats des calculs intermédiaires.
Fonctionnement d'un TPU
Cloud TPU a été conçu par Google comme un processeur matriciel spécialisé dans les charges de travail des réseaux de neurones. Les TPU ne peuvent pas exécuter de traitement de texte, contrôler les moteurs de fusées ni exécuter de transactions bancaires, mais ils peuvent gérer des opérations matricielles massives dans les réseaux de neurones à des vitesses élevées.
La tâche principale des TPU est le traitement matriciel, qui est une combinaison d'opérations de multiplication et d'addition. Les TPU contiennent des milliers d'accumulateurs multiplicateurs qui sont directement connectés les uns aux autres pour former une grande matrice physique. C'est ce qu'on appelle une architecture de tableau systolique. Cloud TPU v3 contient deux tableaux systoliques de 128 x 128 ALU sur un seul processeur.
L'hôte TPU diffuse les données dans une file d'attente d'ingestion. Le TPU charge les données de la file d'attente d'ingestion et les stocke dans la mémoire HBM. Une fois le calcul terminé, le TPU charge les résultats dans la file d'attente de sortie. L'hôte TPU lit ensuite les résultats de la file d'attente de sortie et les stocke dans la mémoire de l'hôte.
Pour effectuer les opérations matricielles, le TPU charge les paramètres de la mémoire HBM dans l'unité de multiplication matricielle (MXU).
Le TPU charge ensuite les données à partir de la mémoire HBM. À chaque exécution d'une multiplication, le résultat est transmis à l'accumulateur suivant. La sortie correspond à la somme de tous les résultats de multiplication entre les données et les paramètres. Aucun accès à la mémoire n'est requis lors du processus de multiplication matricielle.
Par conséquent, les TPU peuvent atteindre un débit de calcul élevé sur les calculs de réseau de neurones.
Architecture du système TPU
Les sections suivantes décrivent les concepts clés d'un système TPU. Pour en savoir plus sur les termes courants du machine learning, consultez le Glossaire du machine learning.
Si vous débutez avec Cloud TPU, consultez la page d'accueil de la documentation TPU.
Puce TPU
Une puce TPU contient un ou plusieurs TensorCores. Le nombre de TensorCores dépend de la version de la puce TPU. Chaque TensorCore se compose d'une ou plusieurs unités de multiplication de matrices (MXU), d'une unité vectorielle et d'une unité scalaire. Pour en savoir plus sur les TensorCores, consultez A Domain-Specific Supercomputer for Training Deep Neural Networks (Super-ordinateur spécifique au domaine pour l'entraînement de réseaux de neurones profonds).
Une unité matricielle est composée de 256 x 256 (TPU v6e) ou 128 x 128 (versions de TPU antérieures à la version v6e) d'accumulateurs de multiplication dans un tableau systolique. Les unités matricielles fournissent l'essentiel de la puissance de calcul d'un TensorCore. Chaque unité matricielle est capable d'effectuer 16 000 opérations multiply-accumulate par cycle. Toutes les multiplications utilisent des entrées bfloat16, mais toutes les accumulations sont effectuées au format numérique FP32.
L'unité vectorielle est utilisée pour les calculs généraux tels que les activations et les 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.
Pod TPU
Un pod TPU est un ensemble contigu de TPU regroupés sur un réseau spécialisé. Le nombre de puces TPU dans un pod TPU dépend de la version du TPU.
Tranche
Une tranche est un ensemble de puces situées dans le même pod TPU et connectées par des interconnexions entre puces (ICI) à haut débit. Les tranches sont décrites en termes de puces ou de TensorCores, en fonction de la version du TPU.
La forme de chip et la topologie de chip font également référence aux formes de tranche.
Multislice ou tranche unique
Multislice est un groupe de tranches, qui étend la connectivité TPU au-delà des connexions d'interconnexion entre puces (ICI) et exploite le réseau de centre de données (DCN) pour transmettre des données au-delà d'une tranche. Les données de chaque tranche sont toujours transmises par l'ICI. Grâce à cette connectivité hybride, Multislice permet le parallélisme entre les tranches et vous permet d'utiliser un plus grand nombre de cœurs TPU pour une seule tâche que ce qu'une seule tranche peut accueillir.
Les TPU peuvent être utilisés pour exécuter une tâche sur une seule tranche ou sur plusieurs tranches. Pour en savoir plus, consultez la section Présentation de la multislice.
Cube TPU
Topologie 4x4x4 de puces TPU interconnectées. Cela ne s'applique qu'aux topologies 3D (à partir de la version 4 de TPU).
SparseCore
Les SparseCores sont des processeurs de flux de données qui accélèrent les modèles en s'appuyant sur les représentations vectorielles continues qu'on peut trouver dans les modèles de recommandation. La version v5p comprend quatre SparseCores par puce, et la version v6e deux SparseCores par puce.
Résilience des ICI Cloud TPU
La résilience ICI permet d'améliorer la tolérance aux pannes des liaisons optiques et des commutateurs de circuits optiques (OCS) qui connectent les TPU entre les cubes. (les connexions ICI au sein d'un cube utilisent des liaisons en cuivre qui ne sont pas affectées). La résilience ICI permet de contourner les défaillances de l'OCS et de l'ICI optique. Par conséquent, il améliore la disponibilité de planification des tranches de TPU, avec le compromis d'une dégradation temporaire des performances de l'ICI.
Pour Cloud TPU v4 et v5p, la résilience ICI est activée par défaut pour les tranches d'un cube ou plus, par exemple:
- v5p-128 lorsque vous spécifiez le type d'accélérateur
- 4x4x4 lors de la spécification de la configuration de l'accélérateur
Versions de TPU
L'architecture exacte d'une puce TPU dépend de la version de TPU que vous utilisez. Chaque version de TPU est également compatible avec différentes tailles et configurations de tranches. Pour en savoir plus sur l'architecture du système et les configurations compatibles, consultez les pages suivantes:
Architecture cloud TPU
Google Cloud met à disposition les TPU en tant que ressources de calcul via des VM TPU. Vous pouvez utiliser directement des VM TPU pour vos charges de travail ou les utiliser via Google Kubernetes Engine ou Vertex AI. Les sections suivantes décrivent les principaux composants de l'architecture cloud TPU.
Architecture avec une VM TPU
L'architecture de la VM TPU vous permet de vous connecter directement à la VM physiquement connectée à l'appareil TPU à l'aide de SSH. Une VM TPU, également appelée "worker", est une machine virtuelle exécutant Linux qui a accès aux TPU sous-jacents. 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 du compilateur et de l'environnement d'exécution, ainsi qu'aux messages d'erreur.
Hôte unique, multi-hôte et sous-hôte
Un hôte TPU est une VM qui s'exécute sur un ordinateur physique connecté à du matériel TPU. Les charges de travail TPU peuvent utiliser un ou plusieurs hôtes.
Une charge de travail à hôte unique est limitée à une seule VM TPU. Une charge de travail multi-hôte distribue l'entraînement sur plusieurs VM TPU. Une charge de travail de sous-hôte n'utilise pas tous les puces d'une VM TPU.
Architecture de nœud TPU (obsolète)
L'architecture du 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, ce qui rend le débogage de l'entraînement et des erreurs TPU difficile.
Passer d'une architecture avec un nœud TPU à une architecture avec une VM TPU
Si vous disposez de TPU utilisant l'architecture de nœud TPU, suivez la procédure ci-dessous pour les identifier, les supprimer et les provisionner à nouveau en tant que VM TPU.
Accédez à la page "TPU" :
Recherchez votre TPU et son architecture sous l'en-tête Architecture. Si l'architecture est "VM TPU", aucune action n'est requise de votre part. Si l'architecture est "Nœud TPU", vous devez supprimer et réprovisionner le TPU.
Supprimez et réprovisionnez le TPU.
Consultez la section Gérer les TPU pour savoir comment supprimer et provisionner à nouveau des TPU.