Architecture du système

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

Aperçu

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. Leur conception repose sur la vaste expérience de Google et sur son leadership en matière de machine learning.

Vous pouvez utiliser Cloud TPU et TensorFlow pour exécuter vos propres 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 à créer des clusters de calcul TensorFlow capables d'exploiter des processeurs, des GPU et des TPU. Les API TensorFlow de haut niveau facilitent l'exécution de modèles dupliqués sur le matériel Cloud TPU.

Vos applications TensorFlow peuvent accéder aux nœuds TPU à partir de conteneurs, d'instances ou de services sur Google Cloud. L'application nécessite une connexion au nœud TPU via votre réseau VPC.

Versions de TPU

Chaque version de TPU définit les caractéristiques matérielles spécifiques d'un appareil 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 appareil TPU et les interfaces réseau disponibles pour la communication entre appareils. Par exemple, chaque version de TPU présente les caractéristiques suivantes :

  • TPU v2 :
    • 8 Gio de mémoire HBM par cœur de TPU
    • Une unité matricielle par cœur de TPU
    • Jusqu'à 512 cœurs de TPU et 4 Tio de mémoire totale dans un pod TPU
  • TPU v3 :
    • 16 Gio de mémoire HBM par cœur de TPU
    • Deux unités matricielles par cœur de TPU
    • Jusqu'à 2 048 cœurs de TPU et 32 Tio de mémoire totale dans un pod TPU

image

Chaque cœur de TPU comprend des unités scalaires, vectorielles et matricielles (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 par cycle. Tandis que les entrées et sorties des unités matricielles sont des valeurs à virgule flottante 32 bits, les unités matricielles effectuent des multiplications 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'un appareil TPU peut effectuer des calculs utilisateur (opérations XLA) indépendamment. Des interconnexions à bande passante élevée permettent aux puces de communiquer entre elles directement sur l'appareil 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 vos charges de travail de machine learning.

Avantages du TPU v3 en termes de performances

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 certains pas 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.

Vous pouvez déterminer si les performances de votre modèle vont s'améliorer à partir d'un TPU v3 en effectuant des analyses comparatives sur différentes versions de TPU et en surveillant les performances à l'aide des outils TensorBoard.

Configurations TPU

Dans un centre de données Google, les appareils TPU sont disponibles dans les configurations suivantes pour Cloud TPU v2 et Cloud TPU v3 :

  • Appareils TPU uniques : appareils TPU individuels qui ne sont pas interconnectés sur un réseau haut débit dédié. Il n'est pas possible de combiner plusieurs types d'appareils TPU uniques en vue de collaborer sur une charge de travail unique.
  • Pods TPU : clusters d'appareils TPU interconnectés sur des réseaux haut débit dédiés.

Appareil TPU unique

Dans un centre de données Google, une configuration d'appareil TPU unique consiste en seul appareil TPU sans connexion réseau haut débit dédiée à d'autres appareils TPU. Votre nœud TPU n'est relié qu'à cet appareil :

image

Les puces des appareils TPU uniques sont interconnectées. Par conséquent, la communication entre les puces ne nécessite pas de ressources processeur hôtes ni de ressources réseau hôtes.

Lorsque vous créez un nœud TPU, vous spécifiez le type de TPU. Par exemple, vous pouvez spécifier v2-8 ou v3-8 pour configurer votre nœud TPU avec un seul appareil. Les appareils TPU uniques ne font pas partie des configurations de pod TPU et n'occupent aucune section d'un pod TPU. Consultez la page Types de TPU pour connaître les configurations des appareils TPU uniques disponibles pour vos nœuds TPU.

Pods TPU

Dans un centre de données Google, une configuration de pod TPU comprend plusieurs appareils TPU interconnectés sur une connexion réseau haut débit dédiée. Les hôtes de votre nœud TPU distribuent vos charges de travail de machine learning sur l'ensemble des appareils TPU.

image

Les puces d'un pod TPU sont interconnectées sur l'appareil. Par conséquent, la communication entre les puces ne nécessite pas de ressources processeur hôtes ni de ressources réseau hôtes. En outre, tous les appareils TPU d'un pod TPU sont interconnectés sur des réseaux haut débit dédiés qui ne nécessitent pas non plus de ressources processeur hôtes ni de ressources réseau hôtes.

Lors de la création d'un nœud TPU, vous indiquez que vous souhaitez un type de TPU qui occupe soit la totalité du pod TPU, soit une fraction de ce pod. Par exemple, un type de TPU v2-512 occupe un pod TPU v2 complet et un type de TPU v2-128 occupe seulement le quart d'un pod TPU v2. Consultez la page Types de TPU pour connaître les configurations de pod de TPU disponibles pour vos nœuds TPU.

image

Le pod TPU v2 propose une configuration maximale de 64 appareils pour un total de 512 cœurs de TPU v2 et 4 Tio de mémoire TPU.

image

Le pod TPU v3 propose une configuration maximale de 256 appareils pour un total de 2 048 cœurs de TPU v3 et 32 Tio de mémoire TPU.

Consultez la section Versions de TPU pour en savoir plus sur les différences d'architecture entre les versions de TPU.

Vous pouvez utiliser l'API Cloud TPU pour automatiser la gestion des TPU pour vos nœuds TPU, quelle que soit leur taille. 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 machine learning. En pratique, la pile logicielle Cloud TPU élimine la complexité au niveau de la création, l'exécution et l'alimentation des programmes Cloud TPU.

Architecture logicielle

Lorsque vous exécutez votre application, TensorFlow génère un graphe de calcul et l'envoie à un nœud TPU via gRPC. Le type de TPU que vous sélectionnez pour votre nœud TPU détermine le nombre d'appareils disponibles pour votre charge de travail. Le nœud TPU compile le graphe de calcul à la volée et envoie le binaire du programme à un ou plusieurs appareils TPU pour exécution. Les entrées du modèle sont souvent stockées dans Cloud Storage. Le nœud TPU transmet les entrées à un ou plusieurs appareils TPU pour qu'elles soient traitées.

Le diagramme ci-dessous illustre l'architecture logicielle Cloud TPU constituée du modèle de réseau de neurones (TPU Estimator), du client TensorFlow, du serveur TensorFlow et du compilateur XLA.

image

TPU Estimator

Les TPU Estimators sont un ensemble d'API de haut niveau qui s'appuient sur des estimateurs simplifiant la création de modèles pour Cloud TPU et tirant parti des performances maximales des TPU. Lorsque vous écrivez un modèle de réseau de neurones utilisant Cloud TPU, il est recommandé d'utiliser les API TPU Estimator.

Client TensorFlow

Les TPU Estimators traduisent vos programmes en opérations TensorFlow, qui sont ensuite converties en graphe de calcul par un client TensorFlow. Ce dernier transmet ensuite le graphe de calcul à un serveur TensorFlow.

Serveur TensorFlow

Un serveur TensorFlow s'exécute sur un serveur Cloud TPU. Lorsque le serveur reçoit un graphe de calcul du client TensorFlow, il réalise les actions suivantes :

  1. Il charge les entrées depuis Cloud Storage.
  2. Il partitionne le graphe en sections pouvant s'exécuter sur un Cloud TPU et en sections devant s'exécuter sur un processeur.
  3. Il génère les opérations XLA correspondant au sous-graphe à exécuter sur Cloud TPU.
  4. Il appelle le compilateur XLA.

Compilateur XLA

XLA est un compilateur à la volée qui prend en entrée les opérations HLO (High Level Optimizer) générées par le serveur TensorFlow. XLA génère un code binaire à exécuter sur le TPU incluant l'orchestration des données depuis la mémoire intégrée vers les unités d'exécution matérielles et la communication entre puces. Le fichier binaire généré est chargé sur le TPU grâce à la connectivité PCIe entre le serveur Cloud TPU et le TPU proprement dit, puis il est lancé pour exécution.

Étape suivante