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 Units) sont des circuits intégrés spécifiques aux applications (Application-Specific Integrated Circuit ou ASIC), développés spécifiquement par Google et permettant d'accélérer les charges de travail de machine learning (ML). Leur conception repose sur la vaste expérience de Google et sur son leadership en matière de machine learning.

Une seule carte TPU contient quatre puces TPU. Chaque puce contient deux cœurs TPU. Les ressources disponibles sur un cœur TPU varient selon la version. Chaque cœur de TPU comprend des unités de multiplication, scalaire et de multiplication matricielle (MXU). Les unités matricielles fournissent l'essentiel de la puissance de calcul d'une puce TPU. Chaque unité matricielle est capable d'effectuer 16 000 opérations multiply-accumulate dans chaque cycle avec une précision réduite de bfloat16. bfloat16 est une représentation en virgule flottante sur 16 bits qui améliore l'entraînement et la précision du modèle par rapport à la représentation à demi-précision de la norme IEEE. Chacun des cœurs d'une carte TPU peut effectuer des calculs utilisateur de manière indépendante. Des interconnexions à bande passante élevée permettent aux puces de communiquer entre elles directement. Des informations techniques plus détaillées sont disponibles dans la section Superordinateur spécifique à un domaine pour l'entraînement des réseaux de neurones profonds.

Les TPU sont conçus pour être adaptés à un pod TPU. Un pod TPU est un superordinateur pouvant comporter 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 relient plusieurs appareils TPU pour fournir davantage 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 hautes performances basée sur le processeur, pour les opérations telles que le chargement et le prétraitement des données en flux continu sur les TPU.

Cloud TPU

Cloud TPU est un service qui vous permet d'accéder 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é maximales 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. De cette manière, vous pouvez facilement déployer d'énormes clusters de calcul, exécuter vos charges de travail, puis réduire ces clusters une fois les charges de travail terminées. L'assistance matérielle intégrée aux puces permet de bénéficier d'un scaling linéaire et performant pour une large gamme de charges de travail de deep learning.

Configurations TPU

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

  • Un pod TPU
  • Une tranche de TPU
  • Un tableau TPU

Pods TPU

Toutes les puces d'un pod TPU sont connectées directement les unes aux autres via une interconnexion haut débit, ce qui élimine les retards de communication dans les hôtes de processeur. Les puces sont connectées dans un torchon 2D, chaque puce communiquant directement à quatre voisins. Cette architecture permet d'obtenir d'excellentes performances pour des modèles de communication courants dans des charges de travail de ML telles que la réduction totale. Les détails de l'architecture et des caractéristiques de performances de Cloud TPU v2 et v3 sont présentés dans le tableau suivant, ainsi que des informations publiées sur la page Un superordinateur propre à l'entraînement d'un réseau de neurones profond.

Caractéristique TPU v2 TPU v3
Liens réseau x Gbit/s/Chip 4 x 496 4 x 656
Terabits de bande passante bibande 15,9 42

Tranches TPU

Une tranche de TPU est une partie d'un pod TPU. Si vous n'avez pas besoin des ressources d'un pod entier, vous pouvez utiliser une partie d'un pod. Plusieurs configurations de tranche sont disponibles pour en savoir plus. Consultez la section Tarifs de Cloud TPU dans Cloud TPU.

Tableau TPU unique

Une configuration TPU monocarte est une carte autonome avec 4 puces TPU (8 cœurs TPU) sans connexion réseau à d'autres cartes TPU. Les appareils TPU uniques ne font pas partie d'une configuration de pod TPU et n'occupent aucune section d'un pod TPU. Consultez la page "Types de TPU" pour connaître les configurations de TPU uniques.

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) des cœurs, les interconnexions entre les cœurs de chaque carte TPU et les interfaces réseau disponibles pour{101. }communication entre appareils. Les versions de TPU disponibles sont les suivantes: v1 et v2.

TPU v2

Une carte TPU v2 contient quatre puces TPU ayant chacune deux cœurs. Il existe 8 Gio de mémoire HBM par cœur de TPU et un MXU par cœur. Un pod TPU v2 dispose de jusqu'à 512 cœurs de TPU et 4 Tio de mémoire.

image

TPU v3

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

image

Avantages 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 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 résiduels (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 sur les performances du pipeline peut vous aider à résoudre les problèmes de flux d'entrée.

Architecture de VM Cloud TPU

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

image

Un pod TPU contient un hôte TPU pour chaque carte TPU.

image

Votre mode d'interaction avec l'hôte TPU (et la carte TPU) dépend de l'architecture des VM TPU que vous utilisez: les nœuds TPU ou les VM TPU.

Nœuds TPU

Les nœuds TPU sont 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'y a 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 root à la machine pour pouvoir exécuter le code de votre choix. Vous pouvez accéder aux journaux de débogage et aux messages d'erreur directement depuis le compilateur et l'environnement d'exécution du 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

Étant donné qu'il n'y a pas de VM utilisateur, vous n'avez pas besoin de réseau, de cloud privé virtuel ni de pare-feu entre votre code et la VM TPU, améliorant les performances de votre pipeline d'entrée. Les VM TPU sont moins chères, car vous n'avez pas à payer 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 qui est présente sur chaque VM TPU. Cette bibliothèque comprend le compilateur XLA utilisé pour compiler des 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 d'exécuter votre code Python s'exécutant sur une VM utilisateur, vous pouvez l'exécuter directement sur l'hôte TPU.

image

Pour plus d'informations sur TensorFlow et Cloud TPU, consultez la page Exécuter des modèles TensorFlow sur Cloud TPU.

Nœuds TPU

L'architecture du système de nœuds Cloud TPU a été créée à l'origine 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. Elles n'exécutent pas Python ou aucun code utilisateur non représenté sous la forme d'un 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 graphiques XLA et des instructions d'exécution via des connexions gRPC TensorFlow TensorFlow et de les exécuter sur les serveurs TensorFlow. Une VM utilisateur est requise pour chaque hôte TPU.

image

JAX

VM TPU

Les VM TPU n'ont pas besoin de VM utilisateur, car vous pouvez exécuter votre code directement sur les hôtes TPU.

image

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

Nœuds TPU

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

image

Étape suivante