TPU7x (Ironwood)

Cette page décrit l'architecture et les configurations disponibles pour TPU7x, le dernier TPU disponible sur Google Cloud. TPU7x est la première version de la famille Ironwood, Google Cloudle TPU de septième génération. La génération Ironwood est conçue pour l'entraînement et l'inférence d'IA à grande échelle.

Avec une empreinte de 9 216 puces par pod, TPU7x présente de nombreuses similitudes avec TPU v5p. TPU7x offre des performances élevées pour les modèles denses et MoE à grande échelle, le pré-entraînement, l'échantillonnage et l'inférence intensive en décodage.

Pour utiliser TPU7x, vous devez utiliser Google Kubernetes Engine (GKE). Pour en savoir plus, consultez À propos des TPU dans GKE.

Vous pouvez également utiliser TPU7x et GKE avec TPU Cluster Director. TPU Cluster Director est disponible via une réservation en mode "All Capacity" (Capacité totale), qui vous donne un accès complet à toute votre capacité réservée (sans retenue) et une visibilité totale sur la topologie du matériel TPU, l'état d'utilisation et l'état de santé. Pour en savoir plus, consultez Présentation du mode "All Capacity" (Capacité totale).

Pour accéder à TPU7x, contactez votre équipe chargée du compte.

Architecture du système

Chaque puce TPU7x contient deux TensorCores et quatre SparseCores. Le tableau suivant présente les principales spécifications et leurs valeurs pour TPU7x par rapport aux générations précédentes.

Spécification v5p v6e (Trillium) TPU7x (Ironwood)
Nombre de puces par pod 8960 256 9216
Puissance de calcul maximale par puce (BF16) (TFLOPs) 459 918 2307
Puissance de calcul maximale par puce (FP8) (TFLOPs) 459 918 4614
Capacité de mémoire HBM par puce (Gio) 95 32 192
Bande passante HBM par puce (Gio/s) 2575 1638 7380
Nombre de processeurs virtuels (VM à quatre puces) 208 180 224
RAM (Go) (VM à quatre puces) 448 720 960
Nombre de TensorCores par puce 2 1 2
Nombre de SparseCores par puce 4 2 4
Bande passante d'interconnexion bidirectionnelle entre puces (ICI) par puce (Go/s) 1200 800 1200
Bande passante du réseau de centre de données (DCN) par puce (Gbit/s) 50 100 100

Le schéma suivant illustre l'architecture d'Ironwood :

Schéma de l'architecture Ironwood

Hiérarchie de la mémoire

TPU7x est doté d'un système de mémoire à plusieurs niveaux. La gestion du déplacement des données entre ces niveaux est essentielle pour les performances :

  • Mémoire à haut débit (HBM) : chaque puce est équipée de 192 Go de mémoire HBM, avec une bande passante d’environ 7,37 To/s. La grande capacité de mémoire HBM permet d'utiliser des tailles de lot importantes, ce qui peut améliorer le débit. Toutefois, malgré sa taille, la mémoire HBM peut toujours être un goulot d'étranglement, en particulier pour les opérations vectorielles liées à la mémoire ou les modèles d'accès aux données inefficaces.
  • Mémoire vectorielle (VMEM) : la mémoire VMEM est une mémoire SRAM sur puce plus petite (mémoire statique à accès aléatoire) avec une bande passante nettement supérieure à celle de la mémoire HBM pour l'unité de multiplication matricielle (MXU). Cette mémoire sert de bloc-notes à haute vitesse pour les noyaux personnalisés. La taille de ce tampon est un paramètre réglable. L'optimisation de la taille de la mémoire tampon est essentielle pour régler les noyaux Pallas personnalisés, car leurs tailles de bloc sont souvent limitées par la mémoire VMEM disponible.
  • Mémoire hôte et PCIe : chaque ensemble de quatre puces TPU est connecté à un hôte de processeur à l'aide d'un réseau PCIe. Bien que cette connexion ait une bande passante beaucoup plus faible que la mémoire HBM, la mémoire principale de l'hôte peut être utilisée pour décharger les activations ou les états d'optimiseur afin de libérer de la mémoire HBM. Cette technique est particulièrement utile pour gérer la pression sur la mémoire dans les grands modèles.

Pour en savoir plus sur la gestion efficace du déplacement des données entre les niveaux de la hiérarchie de mémoire TPU7x, consultez Optimisations des performances d'Ironwood.

Architecture à double chiplet

Le modèle de programmation Ironwood vous permet d'accéder à deux chiplets TPU au lieu d'une seule architecture de cœur logique (également appelée MegaCore) utilisée dans les générations précédentes (TPU v4 et v5p). Cette modification améliore la rentabilité et l'efficacité de la fabrication de la puce. Bien qu'il s'agisse d'un changement architectural, la nouvelle conception vous permet de réutiliser les modèles logiciels existants avec un minimum de modifications.

Les TPU Ironwood sont composés de deux chiplets distincts, chacun avec son propre espace mémoire dédié. Cela diffère de l'espace mémoire unifié de l'architecture MegaCore.

  • Composition du chiplet : chaque chiplet est une unité autonome avec un TensorCore, deux SparseCores et 96 Go de mémoire à haut débit (HBM).

  • Interconnexion à haut débit : les deux chiplets sont connectés par une interface die-to-die (D2D) six fois plus rapide qu'une liaison d'interconnexion entre puces (ICI) 1D. La communication entre les chiplets est gérée à l'aide d'opérations collectives.

Modèle de programmation et exposition du framework

Le modèle de programmation d'Ironwood est semblable à celui des générations de TPU antérieures à la version v4, telles que TPU v3. La nouvelle architecture est exposée de la manière suivante :

  • Deux appareils par puce : les frameworks tels que JAX exposent chaque puce Ironwood en tant que deux "appareils" distincts, un pour chaque chiplet.

  • Spécification du chiplet : vous pouvez spécifier le chiplet à utiliser pour un calcul. JAX ajoute une quatrième dimension à la spécification de la topologie pour distinguer les chiplets. Cette conception vous permet de réutiliser les modèles logiciels existants avec un minimum de modifications.

Pour en savoir plus sur l'optimisation des performances avec l'architecture à double chiplet, consultez Recommandations de performances pour l'architecture à double chiplet d'Ironwood

Configurations compatibles

Les puces TPU7x sont directement connectées aux puces voisines les plus proches en trois dimensions, ce qui crée un maillage 3D de connexions réseau. Les tranches de plus de 64 puces sont constituées d'un ou de plusieurs "cubes" 4x4x4 de puces.

Les puces TPU7x ont une topologie d'interconnexion en tore 3D. Cette topologie permet aux tranches d'évoluer jusqu'à 9 216 puces. Elle dispose d'une bande passante bidirectionnelle de 200 Go/s par axe pour la communication entre les puces d'un pod.

Le tableau suivant présente les formes de tranche 3D courantes compatibles avec TPU7x :

Topologie Puces TPU Hôtes VM Cubes Champ d'application
2x2x1 4 1 1 1/16 Un seul hôte
2x2x2 8 2 2 1/8 Multi-hôtes
2x2x4 16 4 4 1/4 Multi-hôtes
2x4x4 32 8 8 1/2 Multi-hôtes
4x4x4 64 16 16 1 Multi-hôtes
4x4x8 128 32 32 2 Multi-hôtes
4x8x8 256 64 64 4 Multi-hôtes
8x8x8 512 128 128 8 Multi-hôtes
8x8x16 1024 256 256 16 Multi-hôtes
8x16x16 2048 512 512 32 Multi-hôtes

VM TPU7x

Chaque machine virtuelle (VM) TPU7x contient quatre puces. Chaque VM a accès à deux nœuds NUMA. Pour en savoir plus sur les nœuds NUMA, consultez Non-uniform memory access sur Wikipédia.

Toutes les tranches TPU7x utilisent des VM à hôte complet et à quatre puces. Les caractéristiques techniques d'une VM TPU7x sont les suivantes :

  • Nombre de processeurs virtuels par VM : 224
  • RAM par VM : 960 Go
  • Nombre de nœuds NUMA par VM : 2

Hyperdisk

Par défaut, le disque de démarrage de la VM pour TPU7x est Hyperdisk Balanced. Vous pouvez associer d'autres disques à votre VM TPU pour bénéficier d'un espace de stockage supplémentaire. Les types de disques suivants sont compatibles avec TPU7x :

  • Hyperdisk Balanced
  • Hyperdisk ML

Pour en savoir plus sur Hyperdisk, consultez Présentation d'Hyperdisk. Pour en savoir plus sur les options de stockage pour Cloud TPU, consultez Options de stockage pour les données Cloud TPU.

Étape suivante