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,le TPU de septième génération de Google Cloud. 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 pouvez utiliser Google Kubernetes Engine (GKE) ou Compute Engine. Pour en savoir plus sur l'utilisation des TPU avec GKE, consultez À propos des TPU dans GKE.

Vous pouvez également utiliser TPU7x et GKE avec le mode "Toute capacité". Le mode "Toute la capacité" est disponible via une réservation en mode "Toute la capacité". Il vous donne un accès complet à toute votre capacité réservée (sans retenue) et une visibilité totale sur la topologie matérielle des TPU, l'état d'utilisation et l'état de fonctionnement. Pour en savoir plus, consultez Présentation du mode "Toute capacité".

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 les TPU 7x par rapport aux générations précédentes.

Spécification v5p v6e (Trillium) TPU7x (Ironwood)
Nombre de puces par pod 8 960 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 (Go/s) 2765 1638 7380
Nombre de vCPU (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 bidirectionnelle d'interconnexion entre puces (ICI) par puce (GBps) 1200 800 1200
Bande passante 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

Le TPU7x dispose d'un système de mémoire à plusieurs niveaux. La gestion du transfert de 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é HBM permet d'utiliser des tailles de lot importantes, ce qui peut améliorer le débit. Toutefois, malgré sa taille, la 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 VMEM est une SRAM (mémoire vive statique) intégrée plus petite, avec une bande passante beaucoup plus élevée vers l'unité de multiplication matricielle (MXU) que la HBM. Cette mémoire sert de bloc-notes à grande vitesse pour les noyaux personnalisés. La taille de cette mémoire tampon est un paramètre ajustable. L'optimisation de la taille de la mémoire tampon est essentielle pour ajuster les noyaux Pallas personnalisés, car leurs tailles de bloc sont souvent limitées par la 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 HBM, la mémoire principale de l'hôte peut être utilisée pour décharger les activations ou les états de l'optimiseur afin de libérer de la HBM. Cette technique est particulièrement utile pour gérer la pression de la mémoire dans les grands modèles.

Pour en savoir plus sur la gestion efficace des transferts de 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'un seul cœur logique (également appelé architecture MegaCore) utilisé dans les générations précédentes (TPU v4 et v5p). Ce changement améliore la rentabilité et l'efficacité de la fabrication du chip. Bien que cela représente un changement d'architecture, 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 des chiplets : 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 inter-puces (ICI) 1D. La communication entre les chiplets est gérée à l'aide d'opérations collectives.

Exposition du modèle et du framework de programmation

Le modèle de programmation pour Ironwood est semblable à celui des générations de TPU antérieures à la v4, comme le TPU v3. La nouvelle architecture est exposée des manières suivantes :

  • Deux appareils par puce : les frameworks tels que JAX exposent chaque puce Ironwood sous la forme de deux "appareils" distincts, un pour chaque chiplet.

  • Spécification des chiplets : vous pouvez spécifier le chiplet à utiliser pour un calcul. JAX ajoute une quatrième dimension à la spécification de la topologie pour faire la distinction entre 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'obtention de performances optimales 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 disposent d'une topologie d'interconnexion en tore 3D. Cette topologie permet aux tranches de passer à 9 216 puces. Il dispose d'une bande passante bidirectionnelle de 200 Gbit/s par axe pour la communication entre les puces d'un même 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 Hôte unique
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 1 024 256 256 16 Multi-hôtes
8x16x16 2 048 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. Voici les caractéristiques techniques d'une VM TPU7x :

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

Modèles compatibles

Un modèle compatible est ajusté et validé pour le matériel TPU7x. Vous pouvez déployer et exécuter les modèles compatibles sans modifier le code. Les modèles compatibles sont configurés pour tirer parti des fonctionnalités TPU7x, y compris l'architecture à double chiplet, les interconnexions à haut débit (ICI), la mémoire HBM et le calcul de précision mixte (comme FP8 et BF16).

Google et la communauté Open Source proposent des recettes, des configurations et une assistance pré-optimisées pour les charges de travail d'entraînement et d'inférence sur les TPU7x.

Modèles optimisés pour l'entraînement

Les modèles optimisés pour l'entraînement incluent des recettes préconfigurées et des hyperparamètres adaptés aux frameworks d'entraînement distribué. Ces recettes configurent des stratégies de parallélisme optimales (y compris le parallélisme des données, des tensors, des pipelines et du contexte) pour maximiser le débit et l'efficacité du scaling sur les topologies TPU7x à hôte unique et multislices.

Pour trouver des recettes d'entraînement, des guides et des benchmarks de performances pour les modèles compatibles sur TPU7x, consultez les ressources suivantes :

  • Recettes d'entraînement de TPU pour TPU7x : dépôt GitHub officiel contenant des recettes d'entraînement préconfigurées, des configurations de précision (FP8 et BF16), des benchmarks de temps de traitement et des métriques de débit pour les modèles open source populaires sur TPU7x.
  • Dépôt MaxText : framework d'entraînement LLM Open Source hautes performances écrit en JAX et optimisé pour Cloud TPU.
  • Entraîner un modèle à l'aide de TPU7x : guide pour provisionner des ressources TPU7x et déployer des charges de travail d'entraînement avec MaxText.

Modèles optimisés pour l'inférence

Les modèles optimisés pour l'inférence sont validés pour la mise en service hautes performances à l'aide de vLLM et du plug-in tpu-inference. Ces modèles sont testés pour vérifier leur exactitude, leur précision numérique et leur débit de diffusion, en tirant parti des kernels PagedAttention optimisés et du parallélisme tensoriel sur TPU7x.

Pour trouver des recettes de diffusion et des outils de déploiement pour les modèles compatibles sur TPU7x, consultez les ressources suivantes :

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 obtenir de l'espace de stockage supplémentaire. Les types de disques suivants sont compatibles avec les TPU7x :

  • Volume Hyperdisk équilibré
  • 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.

Étapes suivantes