TPU v4

Ce document décrit l'architecture et les configurations prises en charge de Cloud TPU v4.

Architecture du système

Chaque puce TPU v4 contient deux TensorCores. Chaque TensorCore comporte quatre unités de multiplication matricielle (MXU), une une unité vectorielle et une unité scalaire. Le tableau suivant présente les principales caractéristiques pour un pod TPU v4.

Caractéristiques clés	Valeurs des pods v4
Pic de calcul par puce	275 téraflops (bf16 ou int8)
Capacité et bande passante HBM2	32 Gio, 1 200 Go/s
Puissance minimale/moyenne/maximale mesurée	90/170/192 O
Taille du pod TPU	4 096 puces
Topologie d'interconnexion	Maillage 3D
Pic de calcul par pod	Exaflops 1.1 (bf16 ou int8)
Réduction globale de la bande passante par pod	1,1 Po/s
Bande passante bissectionnelle par pod	24 To/s

Le schéma suivant illustre une puce TPU v4.

Maillage 3D et tore 3D

Les TPU v4 sont directement connectés aux puces voisines les plus proches en 3 ce qui donne un maillage 3D de connexions réseau. Les connexions peut être configurée en tant que tore 3D sur des tranches où la topologie, AxBxC, est : 2A=B=C ou 2A=2B=C, où chaque dimension est un multiple de 4. Par exemple : 4 x 4 x 8, 4 x 8 x 8 ou 12 x 12 x 24. En général, les performances d'un tore 3D sera plus efficace qu'une configuration de maillage 3D. Pour plus d'informations, voir Topologies tori tordues.

Avantages de TPU v4 par rapport à v3

Cette section décrit les avantages de TPU v4 en termes de performances

Système de mémoire

L'architecture NUMA (Non Uniform Memory Access) est une architecture de mémoire informatique destinée aux machines qui ont plusieurs processeurs. Chaque processeur dispose d'un accès direct à un bloc mémoire. Un processeur et sa mémoire sont appelés nœuds NUMA. NUMA nœuds sont connectés à d'autres NUMA nœuds directement adjacents les uns aux autres. Un processeur d'un nœud NUMA peut accéder à la mémoire d'un autre nœud NUMA, mais cet accès est plus lent mémoire dans un nœud NUMA.

Les logiciels exécutés sur une machine à plusieurs processeurs peuvent placer les données nécessaires à un processeur dans son nœud NUMA, ce qui augmente le débit de la mémoire. Pour en savoir plus sur NUMA, consultez la section Accès à la mémoire non uniforme. sur Wikipédia.

Vous pouvez profiter des avantages de la localité NUMA en liant votre script d'entraînement au nœud 0 NUMA.

Pour activer la liaison de nœuds NUMA:

1. Installez l'outil de ligne de commande numactl.

    $ sudo apt-get update
    $ sudo apt-get install numactl
   

2. Liez votre code de script au nœud 0 NUMA.

    $ numactl --cpunodebind=0 python3 your-training-script
   

Activez la liaison de nœuds NUMA si:

• Si votre charge de travail dépend fortement des charges de travail du processeur (par exemple, les images la classification et les charges de travail liées aux recommandations), quel que soit le framework.
• Si vous utilisez une version d'exécution de TPU sans suffixe -pod (par exemple, tpu-vm-tf-2.10.0-v4).

Autres différences du système de mémoire:

• Les puces TPU v4 disposent d'un espace mémoire HBM unifié de 32 Gio sur l'ensemble de la puce. ce qui améliore la coordination entre les deux TensorCores intégrés.
• Amélioration des performances HBM grâce aux dernières normes et débits de mémoire.
• Amélioration du profil de performances DMA avec prise en charge intégrée des hautes performances avec un niveau de précision de 512 milliards.

TensorCores

• Deux fois plus d'unités matricielles et une fréquence d'horloge plus élevée, avec un maximum de 275 TFLOPS.
• Bande passante de transposition et de permutation multipliées par 2
• Modèle d'accès à la mémoire de load-store pour la mémoire commune (Cmem).
• Bande passante plus rapide pour les charges de travail MXU et prise en charge du mode 8 bits pour réduire des tailles de lot et une latence d'inférence améliorée.

Interconnexion entre puces

Six liaisons d'interconnexion par puce pour activer les topologies de réseau différents diamètres du réseau.

Autre

• Interface PCIE 3e génération x16 vers l'hôte (connexion directe)
• Modèle de sécurité amélioré.
• Efficacité énergétique améliorée.

Configurations

Un pod TPU v4 est composé de 4 096 puces interconnectées avec des configurations les liaisons à haut débit. La mise en réseau flexible du TPU v4 vous permet de connecter les puces dans un d'une tranche de pod de même taille de plusieurs manières. Lorsque vous créez une tranche de pod TPU, vous devez spécifier la version de TPU et le nombre de ressources TPU dont vous avez besoin. Lorsque vous créer une tranche de pod TPU v4, vous pouvez spécifier son type et sa taille de deux manières: AcceleratorType et AccleratorConfig.

Remarque :AcceleratorType et AccleratorConfig s'excluent mutuellement. vous pouvez n'en utilisez qu'un lors de la création d'un TPU.

Utiliser AcceleratorType

Utilisez AcceleratorType lorsque vous ne spécifiez pas de topologie. Pour configurer des TPU v4 à l'aide de AcceleratorType, utilisez l'option --accelerator-type lors de la création de votre Tranche de pod TPU. Définissez --accelerator-type sur une chaîne contenant le TPU. et le nombre de TensorCores que vous souhaitez utiliser. Par exemple, pour créer Tranche de pod v4 comportant 32 TensorCores, vous utiliseriez --accelerator-type=v4-32.

La commande suivante crée une tranche de pod TPU v4 avec 512 TensorCores à l'aide de l'indicateur --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v4-512 \
    --version=tpu-vm-tf-2.17.0-pod-pjrt

Le nombre qui suit la version de TPU (v4) spécifie le nombre de TensorCores. Un TPU v4 comporte deux TensorCores. Le nombre de puces TPU serait 512/2 = 256.

Utiliser AcceleratorConfig

Utilisez AcceleratorConfig lorsque vous souhaitez personnaliser la topologie physique de la tranche TPU. Cela est généralement nécessaire pour le réglage des performances avec le pod de tranches de plus de 256 chips.

Pour configurer des TPU v4 à l'aide de AcceleratorConfig, utilisez --version et les Indicateurs --topology. Définissez --version sur la version de TPU que vous souhaitez utiliser, puis --topology à la disposition physique des puces TPU dans la tranche de pod.

Vous spécifiez une topologie de TPU à l'aide d'un 3-tuple, AxBxC où A<=B<=C et A, B, C sont Tous les nombres sont <= 4 ou sont tous des multiples entiers de 4. Les valeurs A, B et C sont le nombre de chips dans chacune des trois dimensions. Par exemple, pour créer un pod v4, avec 16 chips, définissez --version=v4 et --topology=2x2x4.

La commande suivante crée une tranche de pod TPU v4 avec 128 puces TPU organisées en Un tableau 4x4x8:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-vm-tf-2.17.0-pod-pjrt

Les topologies où 2A=B=C ou 2A=2B=C ont également des variantes de topologie. optimisé pour la communication tout-à-tout, par exemple 4 × 4 × 8, 8 × 8 × 16 et 12 × 12 × 24. On parle de topologie toris torsadées.

Les illustrations suivantes présentent certaines topologies TPU v4 courantes.

Les tranches de pod plus volumineuses peuvent être créées à partir d'un ou de plusieurs "cubes" de 4 x 4 x 4 pixels. de chips.

Topologies tori tordues

Dans la version 4, certaines formes de tranches de tore 3D permettent d'utiliser ce que l'on appelle tore torsadée. Par exemple, vous pouvez organiser deux cubes v4 sous la forme d'un bloc 4x4x8 ou 4x4x8_twisted. Les topologies torsadées offrent une bissection la bande passante réseau. L'augmentation de la bande passante bissectionnelle est utile pour les charges de travail qui utilisent des modèles de communication. Les topologies torsadées peuvent améliorer les performances pour la plupart avec des charges de travail de représentation vectorielle continue TPU qui en profitent le plus.

Pour les charges de travail qui utilisent le parallélisme des données comme seule stratégie de parallélisme, la technique torsadée peuvent être légèrement plus performantes. Pour les LLM, les performances à l'aide d'un la topologie peut varier en fonction du type de parallélisme (DP, MP, etc.). Excellent consiste à entraîner votre LLM avec et sans topologie torsadée pour déterminer qui offre les meilleures performances pour votre modèle. Quelques tests Modèle MaxText FSDP a constaté une ou deux améliorations de MFU à l'aide d'une topologie torsadée.

Le principal avantage des topologies torsadées est qu'elles transforment (par exemple, 4×4×8) en topologie symétrique étroitement liée. La topologie symétrique présente de nombreux avantages:

• Équilibrage de charge amélioré
• Bande passante bissectionnelle plus élevée
• Routes de paquets plus courtes

Au final, ces avantages se traduisent par de meilleures performances pour de nombreux des modèles de communication.

Le logiciel TPU accepte des toris torsadés sur les tranches où la taille de chaque tranche est égale ou deux fois supérieure à la plus petite dimension. Pour (par exemple, 4 x 4 x 8, 4 × 8 × 8 ou 12 x 12 x 24).

À titre d'exemple, prenons cette topologie de tore 4×2 avec des TPU étiquetés avec leur (X,Y) dans la tranche:

Pour plus de clarté, les arêtes de ce graphique de topologie sont représentées par des arêtes non orientées. Dans chaque arête est une connexion bidirectionnelle entre les TPU. Nous faisons référence au entre un côté de cette grille et le côté opposé sous forme d'arêtes enveloppantes, comme indiqué dans le diagramme.

En tordant cette topologie, on obtient un tore torsadé 4×2 entièrement symétrique de base:

Tout ce qui a changé entre ce diagramme et le précédent est le Y et des bords enveloppés. Au lieu de se connecter à un autre TPU avec le même X elles ont été décalées pour se connecter au TPU X+2 mod 4.

La même idée se généralise à différentes tailles de dimension et à différents nombres de . Le réseau résultant est symétrique, tant que chaque dimension est égale ou deux fois la taille de la plus petite dimension.

Consultez la section Utiliser AcceleratorConfig pour découvrir comment pour spécifier une configuration tori torsadée lors de la création d'un Cloud TPU.

Le tableau suivant présente les topologies torsadées acceptées et une d'augmentation de la bande passante bissectionnelle grâce à elles par rapport aux topologies non torsadées.

Topologie torsadée	Augmentation théorique de la bissection bande passante par rapport à un tore non torsadé
4×4×8_torsadé	~70%
8x8x16_torsadé
12×12×24_torsadé
4×8×8_torsadé	~40%
8×16×16_torsadé

Remarque :La topologie torsadée n'est utilisée automatiquement que si le suffixe _twisted est spécifié avec l'indicateur --topology. Par exemple, pour spécifier une image torsadée de 4 x 4 x 8 pixels, vous devez spécifier l'option topology comme suit: 4x4x8_twisted.

Variantes de topologie TPU v4

Certaines topologies contenant le même nombre de puces peuvent être disposées en différentes de différentes manières. Par exemple, une tranche de pod TPU avec 512 puces (1 024 TensorCores) peut être configurée selon les topologies suivantes: 4x4x32, 4x8x16 ou 8x8x8. Une tranche de pod TPU avec 2 048 puces (4 096 TensorCores), elle offre encore plus d'options de topologie: 4x4x128, 4 x 8 x 64, 4 x 16 x 32 et 8 x 16 x 16.

La topologie par défaut associée à un nombre de puces donné est semblable à un cube. Cette forme est probablement la meilleure pour l'entraînement ML avec parallélisme des données. Autre peuvent être utiles pour les charges de travail avec plusieurs types de parallélisme (par exemple, le parallélisme des modèles et des données, ou le partitionnement spatial d'une simulation, par exemple). Ces charges de travail sont plus performantes si la topologie correspond au parallélisme utilisé. Par exemple, placer un parallélisme des modèles à quatre voies sur la dimension X et les données à 256 voies sur les dimensions Y et Z correspond à une topologie de 4x16x16.

Les modèles présentant plusieurs dimensions de parallélisme sont plus performants de parallélisme mappées aux dimensions de topologie de TPU. Ce sont généralement de données et de modèles de grands modèles de langage (LLM) parallèles. Par exemple, pour un TPU v4, Tranche de pod avec une topologie 8x16x16, dont les dimensions de la topologie du TPU sont 8, 16 et 16. Il est plus performant d'utiliser le parallélisme de modèles à 8 ou 16 voies (mappé sur l'un des les dimensions de topologie de TPU physique). Un parallélisme des modèles à quatre voies ne serait pas optimal avec cette topologie, puisqu'elle n'est alignée sur aucune des dimensions de la topologie TPU, mais ce serait optimal avec une topologie 4x16x32 sur le même nombre de puces.

Les configurations TPU v4 se composent de deux groupes, ceux dont les topologies sont inférieures à 64 puces (petites topologies) et celles dont les topologies sont supérieures à 64 puces (grandes topologies).

Petites topologies v4

Cloud TPU est compatible avec les tranches de pod TPU v4 inférieures à 64 puces : Un cube 4x4x4. Vous pouvez créer ces petites topologies v4 à l'aide de leur nom basé sur TensorCore (par exemple, v4-32) ou leur topologie (par exemple, 2x2x4):

Nom (basé sur le nombre de TensorCore)	Nombre de chips	Topology
v4-8	4	2x2x1
v4-16	8	2x2x2
v4-32	16	2x2x4
v4-64	32	2x4x4

Topologies v4 volumineuses

Les tranches de pod TPU v4 sont disponibles par incréments de 64 puces, aux formes des multiples de 4 sur les trois dimensions. Les dimensions doivent être au format ordre croissant. Plusieurs exemples sont présentés dans le tableau suivant. Quelques ces topologies sont "personnalisées" qui ne peuvent être créées qu'à l'aide Options --type et --topology, car il existe plusieurs façons d'organiser les chips.

Nom (basé sur le nombre de TensorCore)	Nombre de chips	Topology
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
N/A : vous devez utiliser les options --type et --topology.	256	4x4x16
v4-1024	512	8x8x8
v4-1536	768	8x8x12
v4-2048	1 024	8x8x16
N/A : vous devez utiliser les options --type et --topology.	1 024	4x16x16
v4-4096	2 048	8x16x16
…	…	…