Architecture du système

Ce document décrit l'architecture des composants matériels et logiciels du système Cloud TPU.

Les TPU (Tensor Processing Unit) sont des circuits intégrés (ASIC) propres aux applications et développés par Google pour accélérer les charges de travail de machine learning (ML). Les TPU sont conçus dès le départ grâce à l'expérience approfondie et au leadership de Google en matière de machine learning.

Une carte TPU contient quatre puces de TPU. Chaque puce contient deux cœurs de TPU. Les ressources disponibles sur un cœur TPU varient en fonction de la version. Chaque cœur de TPU possède des unités de multiplication (MXU) scalaires, vectorielles et de matrice. Les unités matricielles fournissent la majeure partie de la puissance de calcul d'une puce TPU. Chaque unité matricielle peut effectuer 16 000 opérations d'accumulation dans chaque cycle avec une précision bfloat16 réduite. bfloat16 est une représentation à virgule flottante de 16 bits qui améliore l'entraînement et la précision du modèle par rapport à la représentation à demi-précision IEEE. Chacun des cœurs d'une carte TPU peut exécuter des calculs utilisateur indépendamment. Les interconnexions à haut débit permettent aux puces de communiquer directement entre elles. Vous trouverez des informations techniques plus détaillées dans la section Super-domain spécifique au domaine pour l'entraînement des réseaux de neurones profonds.

Les TPU ont été conçus pour être adaptés à un pod TPU. Un pod TPU est un superordinateur pouvant contenir jusqu'à 2 048 cœurs de TPU, ce qui vous permet de répartir la charge de traitement sur plusieurs cartes TPU. Dans une configuration de pod TPU, des interfaces réseau haut débit dédiées connectent plusieurs appareils TPU ensemble afin de fournir un plus grand nombre de cœurs de TPU et un pool de mémoire TPU plus important pour votre. Charges de travail de machine learning

Chaque carte TPU est connectée à une machine hôte basée sur le processeur à hautes performances pour le chargement et le prétraitement des données afin d'alimenter les TPU.

Cloud TPU

Cloud TPU est un service qui vous donne accès aux TPU via Google Cloud Platform (GCP). Vous pouvez utiliser Cloud TPU pour exécuter des charges de travail de machine learning sur le matériel accélérateur des TPU de Google. Cloud TPU est conçu pour offrir des performances et une flexibilité optimales afin d'aider les chercheurs, les développeurs et les entreprises à entraîner des charges de travail de ML.

Vous pouvez utiliser l'API Cloud TPU pour automatiser la gestion des TPU pour vos TPU. Ainsi, il est facile d'effectuer un scaling à la hausse vers des clusters de calcul volumineux, d'exécuter vos charges de travail et de les réduire lorsque vos charges de travail sont terminées. La compatibilité matérielle intégrée aux puces permet un scaling linéaire efficace sur un large éventail de charges de travail de deep learning.

Configurations TPU

Cloud TPU permet d'accéder à différentes configurations TPU:

  • Un pod TPU
  • Une tranche TPU
  • Une carte en TPU

Pods TPU

Dans un pod TPU, toutes les puces TPU d'un pod TPU sont connectées directement les unes aux autres via une interconnexion haut débit qui ignore les délais de communication entre les hôtes du processeur. Les puces sont connectées dans un tore 2D, où chaque puce communique directement avec quatre voisins. Cette architecture offre des performances particulièrement élevées pour les modèles de communication courants dans les charges de travail de ML, telles que "tout réduire". Les détails architecturaux et les caractéristiques de performances des TPU v2 et v3 sont présentés dans le tableau suivant. Pour en savoir plus, consultez l'article Super-ordinateur spécifique au domaine pour l'entraînement des réseaux de neurones profonds.

Caractéristique TPU v2 TPU v3
Liens réseau x x/s/s/puce 4 x 496 4 x 656
Bia Bande passante de térabits/pod complet 15,9 42

Tranches TPU

Une tranche TPU fait partie d'un pod TPU. Si vous n'avez pas besoin des ressources d'un pod entier, vous pouvez utiliser une partie de celui-ci. Pour en savoir plus, plusieurs configurations de tranche sont disponibles dans la section "Tarifs de Cloud TPU" dans Cloud TPU.

Carte TPU unique

Une configuration TPU à carte unique est une carte autonome avec quatre puces TPU (8 cœurs TPU) 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 une partie d'un pod TPU. Consultez la page "Types de TPU" pour découvrir les configurations de TPU à carte unique disponibles.

Versions de TPU

La version de TPU définit l'architecture de chaque cœur de TPU, la quantité de mémoire à haut débit (HBM) pour chaque cœur de TPU, les interconnexions entre les cœurs de chaque carte TPU et les interfaces réseau disponibles pour{101 }de la communication entre appareils. Les versions de TPU disponibles sont v2 et v3.

TPU v2

Une carte TPU v2 contient quatre puces TPU avec chacune deux cœurs. Il y a 8 Gio de HBM pour chaque cœur de TPU et un MXU par cœur. Un pod TPU v2 possède jusqu'à 512 cœurs de TPU et 4 Tio de mémoire.

image

TPU v3

Une carte TPU v3 contient quatre puces TPU chacune avec deux cœurs. Il existe 16 Gio de mémoire HBM pour chaque cœur de TPU et deux unités matricielles pour chaque cœur. Un pod TPU v3 possède jusqu'à 2 048 cœurs de TPU et 32 Tio de mémoire.

image

Avantages en termes de performances du TPU v3 par rapport à la version 2

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 liés à la mémoire sur les configurations TPU v2 peuvent ne pas obtenir cette même amélioration des performances s'ils sont également liés à la mémoire sur les configurations TPU v3.

  • Dans les cas où les données ne s'adaptent pas à la mémoire sur les configurations TPU v2, le TPU v3 peut améliorer les performances et réduire le calcul 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, le TPU v3 peut autoriser des ResNets plus profonds et des images plus grandes avec RetinaNet.

Avec les TPU v2 qui sont presque liés ("infeed") parce que les étapes d'entraînement attendent une entrée, ils peuvent également être subordonnés aux entrées avec Cloud TPU v3. Ce guide peut vous aider à résoudre les problèmes de flux.

Architectures de VM Cloud TPU

Chaque carte TPU est connectée physiquement à une machine hôte (hôte TPU).

image

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

image

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

Nœuds TPU

Les nœuds TPU constituent l'expérience TPU d'origine. Elles nécessitent une VM utilisateur supplémentaire qui communique avec l'hôte TPU via gRPC. Il n'existe pas d'accès direct à l'hôte TPU.

image

VM TPU

Lorsque vous utilisez des VM TPU, vous vous connectez directement à une VM Google Compute Engine exécutée sur l'hôte TPU. Vous obtenez un accès racine à la machine afin que vous puissiez exécuter n'importe quel code. Vous pouvez accéder aux journaux de débogage et aux messages d'erreur directement à partir du compilateur et de l'environnement d'exécution TPU. Les VM TPU sont compatibles avec de nouveaux cas d'utilisation qui ne sont pas possibles avec les nœuds TPU. Par exemple, vous pouvez exécuter des opérations personnalisées dans le pipeline d'entrée et utiliser le stockage local.

image

En l'absence de VM utilisateur, vous n'avez pas besoin d'un réseau, d'un cloud privé virtuel ou d'un pare-feu entre votre code et la VM TPU, ce qui améliore les performances de votre pipeline d'entrée. De plus, les VM TPU sont moins chères, car vous ne payez pas pour les VM utilisateur.

Frameworks compatibles

Les frameworks tels que JAX, PyTorch et TensorFlow accèdent aux TPU via une bibliothèque partagée appelée libtpu, présente sur chaque VM TPU. Cette bibliothèque inclut le compilateur XLA utilisé pour compiler les programmes TPU, l'environnement d'exécution 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 de votre code Python s'exécutant sur une VM utilisateur, il 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'architecture du nœud Cloud TPU a été initialement conçue pour TensorFlow et son modèle de programmation distribué. Les hôtes TPU ne sont pas accessibles à l'utilisateur et n'exécutent qu'une copie sans interface graphique d'un serveur TensorFlow. Ils n'exécutent pas Python, ni aucun code utilisateur non représenté sous forme de graphe TensorFlow. Le code utilisateur s'exécute dans une VM distante distincte qui communique avec les hôtes TPU sur le réseau.

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 des 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 sur 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

JAX sur les nœuds Cloud TPU 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

Étape suivante