TPU v4

Ce document décrit l'architecture et les configurations compatibles 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 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 du HBM2	32 Gio, 1 200 Go/s
Puissance minimale/moyenne/maximale mesurée	90/170/192 W
Taille du pod TPU	4 096 puces
Topologie d'interconnexion	Maillage 3D
Pic de calcul par pod	Exaflops 1.1 (bf16 ou int8)
Bande passante entièrement réduite 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 disposent d'une connexion directe aux puces voisines les plus proches en 3 dimensions, ce qui génère un maillage 3D de connexions réseau. Les connexions peuvent être configurées sous la forme d'un torus 3D sur des tranches où la topologie (AxBxC) est soit 2A=B=C, soit 2A=2B=C, où chaque dimension est un multiple de 4. Exemples : 4 x 4 x 8, 4 x 8 x 8 ou 12 x 12 x 24. En général, les performances d'une configuration de tore 3D sont meilleures qu'une configuration de maillage 3D. Pour en savoir plus, consultez la section Topologies tori torsadées.

Avantages des TPU v4 en termes de performances par rapport à ceux de la version v3

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

Système de mémoire

NUMA (Non Uniform Memory Access) est une architecture de mémoire d'ordinateur destinée aux machines dotées de plusieurs processeurs. Chaque processeur dispose d'un accès direct à un bloc de mémoire à haut débit. Un processeur et sa mémoire sont appelés "nœud NUMA". Les nœuds NUMA sont connectés à d'autres nœuds NUMA directement adjacents. 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 que l'accès à la mémoire d'un nœud NUMA.

Les logiciels exécutés sur une machine à plusieurs processeurs peuvent placer les données requises par un processeur dans son nœud NUMA, ce qui augmente le débit de la mémoire. Pour plus d'informations sur NUMA, consultez la section Non Uniform Memory Access (Accès non uniforme à la mémoire) sur Wikipédia.

Vous pouvez bénéficier des avantages de NUMA-Localité 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 le code de votre script au nœud NUMA 0.

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

Activez la liaison de nœuds NUMA dans les cas suivants:

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

Autres différences de 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 permet une meilleure coordination entre les deux TensorCores intégrés.
• Amélioration des performances HBM grâce aux normes de mémoire et aux vitesses les plus récentes.
• Amélioration du profil de performances DMA avec prise en charge intégrée des performances élevées à 512 milliards de niveaux de précision.

TensorCores

• Deux fois plus de MXU et une fréquence d'horloge plus élevée pour obtenir un maximum de 275 TFLOPS.
• Bande passante de transposition et de permutation x2.
• Modèle d'accès à la mémoire load-store pour la mémoire commune (Cmem).
• Bande passante de chargement des pondérations des unités matricielles plus rapide et prise en charge du mode 8 bits pour réduire la taille des lots et améliorer la latence d'inférence

Interconnexion entre puces

Six liens d'interconnexion par puce pour activer les topologies de réseaux ayant des diamètres de réseau plus petits.

Autre

• Interface PCIE gen3 x16 à hôte (connexion directe).
• Amélioration du modèle de sécurité
• Efficacité énergétique améliorée.

Configurations

Un pod TPU v4 est composé de 4 096 puces interconnectées avec des liaisons à haut débit reconfigurables. La mise en réseau flexible de TPU v4 vous permet de connecter les puces 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éez une tranche de pod TPU v4, vous pouvez spécifier son type et sa taille de l'une des deux manières suivantes : AcceleratorType et AccleratorConfig.

Remarque :AcceleratorType et AccleratorConfig s'excluent mutuellement. Vous ne pouvez en utiliser qu'un lorsque vous créez 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 la version de TPU et le nombre de TensorCores que vous souhaitez utiliser. Par exemple, pour créer une tranche de pod v4 avec 32 TensorCores, vous devez utiliser --accelerator-type=v4-32.

La commande suivante crée une tranche de pod TPU v4 comportant 512 TensorCore à l'aide de l'option --accelerator-type:

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

Le numéro situé après la version de TPU (v4) indique le nombre de TensorCores. Un TPU v4 comporte deux TensorCores. Le nombre de puces TPU serait donc de 512/2 = 256.

Utiliser AcceleratorConfig

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

Pour configurer des TPU v4 à l'aide de AcceleratorConfig, utilisez les options --version et --topology. Définissez --version sur la version de TPU que vous souhaitez utiliser et --topology sur 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 et C sont tous des <= 4, ou tous des multiples d'entiers de 4. Les valeurs A, B et C correspondent au nombre de chips dans chacune des trois dimensions. Par exemple, pour créer une tranche de pod v4 comportant 16 chips, vous devez définir --version=v4 et --topology=2x2x4.

La commande suivante crée une tranche de pod TPU v4 comportant 128 puces TPU disposées dans un tableau 4x4x8:

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

Les topologies où 2A=B=C ou 2A=2B=C présentent également des variantes de topologie optimisées pour la communication de type "tout-en-un" (par exemple, 4 × 4 × 8, 8 × 8 × 16 et 12 × 12 × 24). Elles sont connues sous le nom de topologies toris torsadées.

Les illustrations suivantes illustrent certaines topologies de TPU v4 courantes.

Les tranches de pod plus grandes peuvent être construites à partir d'un ou plusieurs "cubes" de chips de 4 x 4 x 4.

Topologies tori torsadées

Certaines formes de tranches de tore 3D v4 permettent d'utiliser ce que l'on appelle une topologie de tore torsadé. Par exemple, deux cubes v4 peuvent être disposés en une tranche de 4x4x8 ou en 4x4x8_twisted. Les topologies torsadées offrent une bande passante bissectionnelle nettement plus élevée. L'augmentation de la bande passante bissectionnelle est utile pour les charges de travail qui utilisent des modèles de communication mondiaux. Les topologies torsadées peuvent améliorer les performances de la plupart des modèles, ce qui est le plus bénéfique pour les charges de travail de représentations vectorielles continues TPU volumineuses.

Pour les charges de travail qui utilisent le parallélisme des données comme seule stratégie de parallélisme, les topologies torsadées peuvent être légèrement meilleures. Pour les LLM, les performances avec une topologie torsadée peuvent varier en fonction du type de parallélisme (DP, MP, etc.). Il est recommandé d'entraîner votre LLM avec et sans topologie torsadée afin de déterminer celle qui offre les meilleures performances pour votre modèle. Certaines expériences menées sur le modèle FSDP MaxText ont permis de constater une à deux améliorations de MFU à l'aide d'une topologie torsadée.

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

• Amélioration de l'équilibrage de charge
• Bande passante bissectionnelle plus élevée
• Routes de paquets plus courtes

Ces avantages se traduisent en fin de compte par de meilleures performances pour de nombreux modèles de communication mondiaux.

Le logiciel TPU accepte les tori torsadés sur les tranches où la taille de chaque dimension est égale ou double à celle de la plus petite dimension. Par exemple, 4 x 4 x 8, 4 × 8 × 8 ou 12 x 12 x 24.

Prenons l'exemple de la topologie de torus 4×2 avec des TPU étiquetés avec leurs coordonnées (X,Y) dans la tranche:

Par souci de clarté, les arêtes de ce graphique de topologie sont représentées comme des arêtes non orientées. En pratique, chaque arête est une connexion bidirectionnelle entre les TPU. Les bords entre un côté de cette grille et le côté opposé sont appelés "arêtes enveloppantes", comme indiqué dans le schéma.

En effectuant une rotation de cette topologie, nous obtenons une topologie de tore torsadée 4×2 complètement symétrique:

Tout ce qui a changé entre ce schéma et le précédent, ce sont les arêtes enveloppantes en Y. Au lieu d'être connectés à un autre TPU ayant la même coordonnée X, ils ont été décalés pour se connecter au TPU ayant les coordonnées X+2 mod 4.

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

Consultez la section Utiliser AcceleratorConfig pour savoir comment 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 compatibles et une augmentation théorique de la bande passante bissectionnelle par rapport aux topologies non torsadées.

Topologie torsadée	Augmentation théorique de la bande passante bissectionnelle par rapport à un tore non torsadé
4×4×8_twistés	~70%
8x8x16_twistés
12×12×24_twistés
4×8×8_twistés	~40%
8×16×16_twisted

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

Variantes de topologie de TPU v4

Certaines topologies contenant le même nombre de chips peuvent être organisées de différentes manières. Par exemple, une tranche de pod TPU comportant 512 puces (1 024 TensorCore) peut être configurée à l'aide des topologies suivantes: 4x4x32, 4x8x16 ou 8x8x8. Une tranche de pod TPU dotée de 2 048 puces (4 096 TensorCores) offre encore plus d'options de topologie: 4x4x128, 4x8x64, 4x16x32 et 8x16x16.

La topologie par défaut associée à un nombre de chips donné est celle qui ressemble le plus à un cube. Cette forme est probablement le meilleur choix pour l'entraînement ML avec parallélisme des données. D'autres topologies peuvent être utiles pour les charges de travail comportant plusieurs types de parallélisme (par exemple, parallélisme des modèles et des données, ou partitionnement spatial d'une simulation). Ces charges de travail fonctionnent mieux si la topologie correspond au parallélisme utilisé. Par exemple, placer le parallélisme des modèles à quatre voies sur la dimension X et le parallélisme des données à 256 voies sur les dimensions Y et Z correspond à une topologie de 4 x 16 x 16.

Les modèles ayant plusieurs dimensions de parallélisme fonctionnent mieux lorsque leurs dimensions de parallélisme sont mappées aux dimensions de topologie de TPU. Il s'agit généralement de grands modèles de langage (LLM) parallèles de données et de modèles. Par exemple, pour une tranche de pod TPU v4 avec une topologie de 8 x 16 x 16, les dimensions de topologie de TPU sont 8, 16 et 16. Il est plus efficace d'utiliser le parallélisme des modèles à 8 ou 16 (mappé sur l'une des dimensions de topologie physique de TPU). Un parallélisme de modèle à quatre voies ne serait pas optimal avec cette topologie, car il n'est aligné sur aucune des dimensions de topologie de TPU, mais il serait optimal avec une topologie de 4x16x32 sur le même nombre de puces.

Les configurations TPU v4 se composent de deux groupes : celles 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 accepte les tranches de pod TPU v4 suivantes de moins de 64 puces, cubiques de 4 x 4 x 4 : Vous pouvez créer ces petites topologies v4 en utilisant soit leur nom basé sur TensorCore (par exemple, v4-32), soit leur topologie (par exemple, 2x2x4):

Nom (basé sur le nombre de TensorCore)	Nombre de puces	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, dont les formes sont des multiples de 4 sur les trois dimensions. Les dimensions doivent être dans l'ordre croissant. Plusieurs exemples sont présentés dans le tableau suivant. Certaines de ces topologies sont des topologies "personnalisées" qui ne peuvent être créées qu'à l'aide des indicateurs --type et --topology, car il existe plusieurs façons d'organiser les chips.

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