- Les TPU sont des puces personnalisées développées par Google pour l’entraînement et l’inférence IA à grande échelle, avec une philosophie de conception différente de celle des GPU
- Ils mettent l’accent sur la scalabilité et l’efficacité énergétique, en concevant conjointement le matériel (par ex. configuration système sur puce, grande mémoire on-chip) et le logiciel (compilateur XLA)
- Leur structure centrale repose sur le systolic array, le pipeline et la compilation anticipée, avec une optimisation pour la plupart des opérations de deep learning, en particulier les multiplications de matrices
- Grâce aux technologies OCI et OCS, il est possible d’obtenir une configuration flexible des nœuds, un traitement parallèle haute performance et un choix varié de topologies
- Les systèmes TPU s’étendent de façon hiérarchique, de la puce unique jusqu’aux très grands multipods, permettant l’entraînement de modèles IA massifs et une utilisation efficace des ressources
Vue d’ensemble et contexte
- Les TPU sont des puces basées sur des ASIC développées par Google pour répondre à la montée en charge massive de ses services IA
- Au départ, Google a hésité entre GPU, FPGA et ASIC, mais l’extension de l’usage du deep learning à partir de 2013, notamment dans la recherche vocale, a accru le besoin de matériel sur mesure
- Aujourd’hui, ils jouent un rôle essentiel dans des services IA majeurs comme Gemini et Veo, et sont largement utilisés pour des modèles de recommandation comme DLRM
Structure d’une puce TPU unique
Configuration de base
- Sur une puce unique de TPUv4, on trouve 2 TensorCore (1 seul pour les TPU spécialisés dans l’inférence)
- Chaque TensorCore est relié à une mémoire CMEM (128MiB) et à une mémoire HBM (32GiB)
Structure interne du TensorCore
- Matrix Multiply Unit (MXU) : basée sur un systolic array 128x128, dédiée aux multiplications de matrices
- Vector Unit (VPU) : exécute les opérations élément par élément générales
- Vector Memory (VMEM; 32MiB) : copie les données depuis la HBM et sert de mémoire de préparation pour le travail
- Scalar Unit + Scalar Memory (SMEM; 10MiB) : gère le flot de contrôle, les opérations scalaires et la gestion des adresses mémoire
Différences structurelles avec les GPU
- Les TPU disposent d’une mémoire on-chip (CMEM, VMEM, SMEM) bien plus importante que les GPU
- La capacité de HBM est plus élevée sur les GPU, et leur nombre de cœurs de calcul est aussi bien supérieur
- Avec TPUv5p, on atteint 500 TFLOPs/sec par puce, et 4.45 ExaFLOPs/sec pour un pod complet (8960 puces)
Philosophie de conception des TPU
1. Systolic array et pipeline
- Dans un systolic array, des éléments de traitement (PE) sont disposés en matrice et transmettent leurs résultats de calcul aux éléments voisins
- Une fois les données injectées, des calculs continus sont possibles sans contrôle supplémentaire, les lectures et écritures mémoire n’ayant lieu qu’en entrée et en sortie
- Cette approche est principalement optimisée pour la multiplication de matrices et les convolutions
- Le traitement en pipeline permet de faire progresser simultanément calcul et déplacement des données afin d’optimiser le débit
Inconvénient du systolic array - la sparsité
- Tous les éléments de calcul restent actifs en permanence, ce qui le rend peu adapté aux matrices creuses
- Si les futurs modèles de DL recherchent une sparsité irrégulière, cela pourrait devenir une limite
2. Compilation anticipée (AoT) et réduction du cache
- Le codesign TPU-XLA réduit au minimum la dépendance au cache, nécessaire pour les accès mémoire irréguliers, ce qui diminue la consommation d’énergie
- Le compilateur XLA analyse à l’avance le graphe d’opérations pour déterminer les schémas d’accès mémoire, et privilégie une exploitation centrée sur la mémoire scratchpad plutôt que sur le cache
- Le
@jit de JAX est une forme intermédiaire entre JIT et AoT : lors de la première exécution, un graphe statique est généré puis compilé en AoT via XLA
- Si la forme des entrées change, une recompilation est nécessaire, et il existe des inefficacités pour le padding dynamique ou les boucles
- Cette approche permet une efficacité énergétique élevée, au prix d’un manque de flexibilité
Efficacité énergétique de TPUv4
- Les puces modernes utilisent de la mémoire HBM3, ce qui réduit la consommation d’énergie
- Comme les opérations mémoire consomment des dizaines à des centaines de fois plus d’énergie que le calcul, minimiser les accès mémoire peut fortement améliorer l’efficacité
Structure multi-puce des TPU
Niveau tray (Tray/Board; 4 puces)
- Un tray se compose de 4 puces TPU (8 TensorCore) et d’un hôte CPU
- La connexion Host↔Chip se fait en PCIe, et la connexion Chip↔Chip via Inter-Core Interconnect (ICI), qui offre une bande passante plus élevée
Niveau rack (Rack; 4x4x4=64 puces)
- Un rack comprend 64 puces TPU, reliées en tore 3D (4x4x4) via ICI et OCS (Optical Circuit Switching)
- Chez Google, on distingue rack, pod et slice
- Rack : unité physique de 64 puces (= cube)
- Pod : unité maximale connectable via ICI et OCS (ex. TPUv4 = 4096 puces = 64 racks)
- Slice : unité abstraite de configuration arbitraire, de 4 puces jusqu’au Superpod
Avantages de l’OCS
- Wraparound : chaque axe forme un anneau (tore 1D), ce qui réduit le nombre maximal de sauts entre nœuds
- Slices non contiguës flexibles : grâce à la structure de commutation OCS, des nœuds physiquement éloignés peuvent être réunis dans une même slice, ce qui facilite l’utilisation des ressources et la maintenance
- Topologies tordues : à nombre de puces identique (x, y, z fixés), il est possible de modifier la structure de connexion pour optimiser la vitesse selon certains schémas de calcul (par ex. twisted torus)
Cas d’usage des topologies
- Forme cubique : adaptée au parallélisme des données et des tenseurs (bande passante maximale)
- Forme linéaire (cigare) : avantageuse pour le parallélisme en pipeline
- Twisted torus : améliore les performances quand une communication all-to-all est nécessaire (par ex. tensor parallel)
Superpod (Full Pod/Superpod; TPUv4 : 4096 puces / 64 racks)
- Plusieurs racks sont reliés via ICI et OCS pour construire un système de très grande taille
- Selon le type de topologie de slice, les performances en bande passante de communication et en traitement parallèle varient
- Grâce à l’OCS, on obtient la flexibilité de slices non contiguës, de topologies tordues, etc.
Niveau multi-pod (Multi-pod/Multislice; TPUv4 : plus de 4096)
- Plusieurs pods sont reliés par le réseau de datacenter (DCN), ce qui fournit une infrastructure d’entraînement à grande échelle, mais avec une bande passante inférieure à celle de l’ICI
- L’entraînement de PaLM a utilisé 2 pods (6144 TPUv4), avec une gestion globale des ressources sur 6 pods
- Lors du développement de modèles à grande échelle, le compilateur XLA optimise les schémas de communication en fonction des paramètres définis par les chercheurs (dimensions de parallélisation, etc.)
- XLA insère les opérations de communication entre slices et pods, ce qui permet un entraînement distribué à grande échelle avec un minimum de modifications du code
Exemples de matériel réel et de schémas
- Rack TPU : un tore 3D 4x4x4 constitue une unité, avec 2 trays (8 puces) par rangée
- Tray TPUv4 : en pratique, il comporte 4 ports PCIe (un par TPU)
- Puce TPUv4 : ASIC au centre, entouré de 4 piles HBM (sur la base de 2 TensorCore)
- Plan de sol de la puce TPUv4i (inférence) : 1 TensorCore, avec une large surface occupée par la CMEM
Conclusion
- Remerciements au Google TPU Research Cloud (TRC) pour son soutien à la recherche
Références
Aucun commentaire pour le moment.