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 d'un pod TPU v4.

Le schéma suivant présente une puce TPU v4.

Pour en savoir plus sur les détails d'architecture et les caractéristiques de performances des TPU v4, consultez TPU v4 : An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings.

Maillage 3D et tore 3D

Les TPU v4 sont directement connectés aux puces voisines les plus proches en trois dimensions, ce qui crée un maillage 3D de connexions réseau. Les connexions peuvent être configurées 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, 4x4x8, 4x8x8 ou 12x12x24. En général, les performances d'une configuration en tore 3D sont meilleures que celles d'une configuration en maillage 3D. Pour en savoir plus, consultez Topologies de tores torsadés.

Avantages en termes de performances des TPU v4 par rapport aux TPU v3

Cette section montre comment exécuter un exemple de script d'entraînement de manière économe en mémoire sur TPU v4, ainsi que les améliorations de performances pour TPU v4 par rapport à TPU v3.

Système de mémoire

L'accès à la mémoire non uniforme (NUMA, Non-Uniform Memory Access) est une architecture de mémoire d'ordinateur pour les machines dotées de plusieurs processeurs. Chaque processeur a un accès direct à un bloc de mémoire à grande vitesse. Un processeur et sa mémoire représentent un 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.

Le logiciel exécuté sur une machine multiprocesseur peut placer les données requises par un processeur dans son nœud NUMA, ce qui augmente le débit de mémoire. Pour en savoir plus sur NUMA, consultez Non uniform memory access sur Wikipédia.

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

Pour activer la liaison des nœuds NUMA :

1. Créez un TPU v4.

2. Connectez-vous au TPU via SSH.

3. Installez l'outil de ligne de commande numactl. numactl vous permet d'exécuter des processus avec une stratégie spécifique de planification ou de placement de mémoire NUMA.

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

4. Liez le code de votre script au nœud NUMA 0. Remplacez your-training-script par le chemin d'accès à votre script d'entraînement.

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

Activez la liaison des nœuds NUMA si :

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

Autres différences entre les systèmes de mémoire :

• Les puces TPU v4 disposent d'un espace de mémoire HBM unifié de 32 Gio sur l'ensemble de la puce, ce qui permet une meilleure coordination entre les deux TensorCores présents sur la puce.
• Amélioration des performances de la mémoire HBM grâce aux dernières normes et vitesses de mémoire.
• Profil de performances DMA amélioré avec prise en charge intégrée des pas hautes performances avec une granularité de 512 octets.

TensorCores

• Deux fois plus d'unités matricielles et une fréquence d'horloge plus élevée, ce qui permet d'atteindre 275 TFLOPS maximum.
• Bande passante de transposition et de permutation doublée.
• Modèle d'accès à la mémoire Load-Store pour la mémoire commune (Cmem).
• Bande passante de chargement des poids MXU plus rapide et prise en charge du mode 8 bits pour permettre des tailles de lot plus petites et une latence d'inférence améliorée.

Interconnexion entre puces

Six liens d'interconnexion par puce pour permettre des topologies de réseau avec des diamètres de réseau plus petits.

Autre

• Interface PCIE Gen3 x16 vers l'hôte (connexion directe).
• Modèle de sécurité amélioré.
• Amélioration de l'efficacité énergétique.

Configurations

Un pod TPU v4 est composé de 4 096 puces qui sont interconnectées par des liens haut débit reconfigurables. La mise en réseau flexible des TPU v4 vous permet de connecter les puces d'une tranche de même taille de plusieurs façons. Lorsque vous créez une tranche TPU, vous devez indiquer la version de TPU et le nombre de ressources TPU dont vous avez besoin. Lorsque vous créez une tranche TPU v4, vous pouvez spécifier son type et sa taille de deux manières : AcceleratorType et AccleratorConfig.

Utiliser AcceleratorType

Utilisez AcceleratorType lorsque vous ne spécifiez pas de topologie. Pour configurer des TPU v4 à l'aide de AcceleratorType, utilisez le flag --accelerator-type lorsque vous créez votre tranche TPU. Définissez --accelerator-type sur une chaîne contenant la version du TPU et le nombre de TensorCores que vous souhaitez utiliser. Par exemple, pour créer une tranche v4 avec 32 TensorCores, vous devez utiliser --accelerator-type=v4-32.

Exécutez la commande gcloud compute tpus tpu-vm create pour créer une tranche TPU v4 avec 512 TensorCores à l'aide du flag --accelerator-type :

  $ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central2-b \
    --accelerator-type=v4-512 \
    --version=tpu-ubuntu2204-base

Le nombre après la version du TPU (v4) spécifie le nombre de TensorCores. Un TPU v4 comporte deux TensorCores. Le nombre de puces TPU est donc de 512/2 = 256.

Pour en savoir plus sur la gestion des TPU, consultez Gérer des TPU. Pour en savoir plus sur l'architecture système de Cloud TPU, consultez Architecture système.

Utiliser AcceleratorConfig

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

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

Vous spécifiez une topologie de TPU à l'aide d'un 3-tuple AxBxC, où A<=B<=C et où A, B et C sont tous <= 4, ou sont tous des multiples entiers de 4. Les valeurs A, B et C correspondent au nombre de puces dans chacune des trois dimensions. Par exemple, pour créer une tranche v4 avec 16 puces, vous devez définir --type=v4 et --topology=2x2x4.

Utilisez la commande gcloud compute tpus tpu-vm create pour créer une tranche TPU v4 avec 128 puces TPU disposées dans un tableau 4x4x8 :

  $ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central2-b \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-ubuntu2204-base

Les topologies où 2A=B=C ou 2A=2B=C ont elles aussi des variantes de topologies optimisées pour la communication all-to-all, par exemple 4×4×8, 8×8×16 et 12×12×24. Il s'agit de topologies en tores torsadés.

Les illustrations suivantes montrent quelques topologies TPU v4 courantes.

Les tranches plus grandes peuvent être constituées d'un ou de plusieurs "cubes" 4x4x4 de puces.

Pour en savoir plus sur la gestion des TPU, consultez Gérer des TPU. Pour en savoir plus sur l'architecture système de Cloud TPU, consultez Architecture système.

Topologies de tores torsadés

Certaines formes de tranches de tore 3D v4 peuvent utiliser ce que l'on appelle une topologie en tore torsadé. Par exemple, deux cubes v4 peuvent être disposés sous la forme d'une tranche 4x4x8 ou 4x4x8_twisted (torsadée). Les topologies torsadées offrent une bande passante bissectionnelle nettement plus élevée. Par exemple, une tranche avec la topologie 4x4x8_twisted offre une augmentation théorique de 70 % de la bande passante bissectionnelle par rapport à une tranche 4x4x8 non torsadée. Une bande passante bissectionnelle accrue est utile pour les charges de travail qui utilisent des modèles de communication globaux. Les topologies torsadées peuvent améliorer les performances de la plupart des modèles, en particulier les grandes charges de travail d'embedding TPU.

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 plus performantes. Pour les LLM, une topologie torsadée peut donner des performances variées selon le type de parallélisme (DP, MP, etc.). La bonne pratique consiste à entraîner votre LLM avec et sans topologie torsadée pour déterminer celle qui offre les meilleures performances pour votre modèle. Certaines expériences sur le modèle FSDP MaxText ont montré une amélioration de la MFU de 1 à 2 en utilisant une topologie torsadée.

Le principal avantage des topologies torsadées est qu'elles transforment une topologie de tore asymétrique (par exemple, 4×4×8) en une 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

Ces avantages se traduisent par une amélioration des performances pour de nombreux modèles de communication mondiaux.

Le logiciel TPU est compatible avec les tores torsadés sur les tranches où la taille de chaque dimension est égale ou deux fois supérieure à la taille de la plus petite dimension. Par exemple, 4x4x8, 4x8x8 ou 12x12x24.

Prenons par exemple cette topologie de tore 4x2 avec des TPU qui sont étiquetés avec leurs coordonnées (X,Y) dans la tranche :

Pour plus de clarté, les arêtes de ce graphique de topologie sont affichées sous forme d'arêtes non orientées. En pratique, chaque arête est une connexion bidirectionnelle entre les TPU. Nous appelons "boucles" les arêtes reliant un côté de cette grille et le côté opposé, comme indiqué dans le diagramme.

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

La seule différence entre ce diagramme et le précédent concerne les boucles de l'axe Y. Au lieu de se relier à un autre TPU avec la même coordonnée X, elles ont été déplacées pour se relier au TPU avec la coordonnée X+2 mod 4.

Le même concept s'applique à différentes tailles de dimensions et à différents nombres de dimensions. Le réseau obtenu est symétrique, à condition que chaque dimension soit égale ou deux fois supérieure à la plus petite dimension.

Consultez Utiliser AcceleratorConfig pour savoir comment spécifier une configuration de tores torsadés lors de la création d'un Cloud TPU.

Le tableau suivant présente les topologies torsadées compatibles et compare l'augmentation théorique de la bande passante bissectionnelle avec topologies torsadées et non torsadées.

Variantes de topologie TPU v4

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

La topologie par défaut associée à un nombre de puces donné est celle qui ressemble le plus à un cube. Cette forme est probablement le meilleur choix pour l'entraînement ML en parallèle des données. D'autres topologies 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). Ces charges de travail fonctionnent mieux si la topologie correspond au parallélisme utilisé. Par exemple, placer le parallélisme de modèle quadridirectionnel sur la dimension X et le parallélisme de données à 256 directions sur les dimensions Y et Z correspond à une topologie 4x16x16.

Les modèles comportant plusieurs dimensions de parallélisme fonctionnent de manière optimale lorsque leurs dimensions de parallélisme sont mappées aux dimensions de la topologie TPU. Il s'agit généralement de grands modèles de langage (LLM) parallèles aux données et aux modèles. Par exemple, pour une tranche de TPU v4 avec une topologie 8x16x16, les dimensions de la topologie TPU sont 8, 16 et 16. De meilleurs performances sont obtenues avec le parallélisme de modèle à 8 ou 16 directions (mappé à l'une des dimensions de la topologie TPU physique). Un parallélisme de modèle à quadridirectionnel serait sous-optimal avec cette topologie, car il n'est aligné sur aucune des dimensions de la topologie TPU. En revanche, il serait optimal avec une topologie 4x16x32 sur le même nombre de puces.

Il existe deux groupes de configurations TPU v4 : celles dont la topologie est inférieure à 64 puces (topologies de petite taille) et celles dont la topologie est supérieure à 64 puces (topologies de grande taille).

Topologies v4 de petite taille

Cloud TPU est compatible avec les tranches TPU v4 de moins de 64 puces, soit un cube 4x4x4. Vous pouvez créer ces topologies v4 de petite taille en utilisant leur nom basé sur le nombre de TensorCores (par exemple, v4-32) ou leur topologie (par exemple, 2x2x4) :

Topologies v4 de grande taille

Les tranches TPU v4 sont disponibles par incréments de 64 puces, avec des formes qui sont des multiples de 4 dans les trois dimensions. Les dimensions doivent être classées par ordre croissant. Vous trouverez plusieurs exemples dans le tableau ci-dessous. Certaines de ces topologies sont des topologies "personnalisées" qui ne peuvent être créées qu'à l'aide des flags --type et --topology, car il existe plusieurs façons d'organiser les puces.