Puces IA : Google va déployer deux architectures très différentes
Il n’y pas un TPU 8, mais un TPU 8t dédié à l’entraînement des modèles d’IA et un TPU 8i pour exécuter les applications qui utilisent ces modèles. Les deux modèles, qui équiperont bientôt les services d’IA de Google Cloud, se différencient par leur architecture réseau.
Google a dévoilé sa huitième génération des puces Tensor Processor Unit (TPU) à l’occasion de son événement annuel Google Cloud Next qui se tenait cette semaine à Las Vegas. En développement depuis deux ans, cette génération devrait arriver graduellement dans les différentes régions du cloud public GCP à partir de la rentrée prochaine.
Cette fois-ci, il existe deux modèles de TPU. Le TPU 8t doit servir aux tâches les plus complexes, comme l’entraînement de LLM. Le TPU 8i motorisera les instances qui exécutent les IA, que ce soit pour faire de l’inférence sur ces LLM, comme pour faire fonctionner les agents qui pilotent ces tâches d’inférence.
« À chaque nouvelle génération de TPU, nous devons anticiper les technologies et les besoins qui existeront au moment de leur mise sur le marché. Il y a plusieurs années, nous avions prévu une augmentation de la demande en matière d’inférence. Et avec l’essor des agents IA, nous avons estimé que les utilisateurs tireraient profit de puces spécialisées soit dans l’entraînement, soit dans les services d’IA », commente Amin Vahdat, le patron des technologies d’IA et d’infrastructure chez Google, lors d’un point presse.
« Dans une telle architecture, il ne suffit pas de redistribuer la charge à une autre puce disponible, il faut surtout éviter que l’erreur d’une puce ne se propage sur la mémoire de toutes les autres. »
Amin VahdatResponsable des technologies d’IA et d’infrastructure, Google
Il insiste pour dire que le TPU 8t n’est pas simplement une version du TPU 8i avec davantage de cœurs de calcul. Ces processeurs auraient des circuits de communication réseau, de mémoire et d’accélération des calculs radicalement différents.
« Par exemple, nous pouvons interconnecter au niveau des accès mémoire 9600 TPU 8t. Ce qui n’est pas le cas des TPU 8i qui communiquent juste en réseau. Cela signifie que nous avons dû mettre dans le TPU 8t des circuits qui, non seulement, garantissent une latence d’accès très faible, mais qui sont aussi extrêmement pointus dans la gestion des erreurs », dit-il.
« Parce que les puces, toutes les puces, plantent régulièrement. Or, dans une telle architecture, il ne suffit pas de redistribuer la charge à une autre puce disponible, il faut surtout éviter que l’erreur d’une puce ne se propage sur la mémoire de toutes les autres », ajoute-t-il.
Deux types de capacités de calcul
De fait la puissance des TPU 8T ne repose véritablement que sur ses capacités de partage des charges de calcul avec ses congénères. Pour le reste, le TPU 8I est individuellement la puce qui a les meilleures caractéristiques techniques. Il embarque 288 Go de mémoire HBM3E pour accéder à des LLM à la vitesse de 8,6 To/s et 384 Mo de mémoire cache afin de conserver plusieurs fenêtres de contexte. Le TPU 8T, lui, embarque 216 Go de mémoire HBM3E, pour stocker des données de travail qu’il lit/écrit à la vitesse de 6,5 To/s, et 128 Mo de cache.
Le TPU 8I embarque des circuits d’accélération MXU qui lui permettent de calculer très rapidement l’endroit, dans le LLM, où se trouvent les données en rapport avec le prompt d’un utilisateur et des circuits XLU qui servent à générer des contenus d’après ces données. Le TPU 8T a, en plus, des circuits SparseCore qui accélèrent la transformation de contenus originaux en vecteurs mathématiques à stocker dans un LLM.
« Le SparseCore est un accélérateur conçu pour les accès mémoire irréguliers qui sont propres aux embeddings. Là où la traditionnelle unité MXU se charge des calculs matriciels, pour déterminer une position, le SparseCore prend en charge les opérations all-gather qui dépendent des données elles-mêmes. En clair, il évite les goulets d’étranglement que sont les opérations nulles, lesquelles affectent les performances sur un processeur classique », explique Diwakar Gupta, ingénieur chez Google Cloud, dans un billet technique.
« [Le SparseCore] évite les goulets d’étranglement que sont les opérations nulles, lesquelles affectent les performances sur un processeur classique. »
Diwakar GuptaIngénieur, Google Cloud
Google indique une puissance de calcul de 10,1 pétaflops avec une précision FP4 sur le TPU 8i contre 12,6 pétaflops avec la même précision pour le TPU 8T. On retiendra surtout que les nouveaux TPU savent désormais gérer des données d’IA avec une précision sur 4 bits, ce qui devrait permettre aux utilisateurs d’obtenir des résultats plus rapidement, pour moins cher. Toutefois, la précision 4 bits est susceptible de générer des réponses un peu moins exactes.
Parmi les autres circuits de calcul présents, les TPU de huitième génération contiennent des cœurs Axion, soit les cœurs ARM de Google. Mais le fournisseur n’indique pas en quelle quantité.
Le processeur Axion, la mémoire cache, deux unités d’accélération MXU et deux unités d’accélération XLU forment le circuit TensorCore. Le TPU 8i possède deux circuits TensorCore. Le TPU 8T a un seul TensorCore et deux circuits SparseCore. En sus, les deux TPU possèdent un circuit ICI qui leur sert de contrôleur pour communiquer en réseau. Ce circuit est totalement différent entre les deux puces.
Deux puces pour deux types de réseau
Les accélérateurs TPU 8T seront interconnectés dans des clusters d’entraînement de modèles d’IA via un nouveau réseau baptisé Virgo. En substance, les liaisons entre les mémoires HBM des TPU 8T sont deux fois plus rapides que sur la génération actuelle. Ces liaisons sont celles d’un cluster de 9600 TPU au maximum, ce que Google appelle un pod. Et les liaisons entre les TPU situés sur différents pods sont quatre fois plus rapides.
Au total, un réseau Virgo permet d’assembler un super cluster qui comprend jusqu’à 134 000 TPU 8T, avec une bande passante cumulée de 47 Pbit/s (47 millions de milliards de bits transmis par seconde). Selon Google, un tel supercluster fournirait une puissance de calcul de 1,6 million d’exaflops (1,6 million de milliards de milliards d’opérations mathématiques à la seconde). Google affirme qu’il peut, en sus, interconnecter de tels superclusters pour grimper jusqu’à plus d’un million de TPU 8T travaillant de concert à l’entraînement d’un modèle.
Pour supporter ce réseau Virgo, le circuit ICI du TPU 8T connecte la mémoire HBM3E interne à six ports logiques externes, individuellement dotés d’une bande passante de 224 Gbit/s.
Google a également implémenté dans ICI des protocoles TPUDirect RDMA et TPUDirect Storage. Ce sont les équivalents maison des protocoles GPUDirect et Spectrum-X de Nvidia. Ils servent à communiquer des paquets sans perte, de sorte qu’ils puissent arriver directement dans une mémoire HBM, sans passer par une étape de décodage au niveau du processeur. Selon Google, ces protocoles permettraient d’accélérer par dix les accès aux baies de stockage.
Rien de tout cela dans le circuit ICI du TPU 8i. S’il dispose bien de six liens de 200 Gbit/s chacun pour les communications de la puce vers le réseau, ce circuit intègre ici un accélérateur CAE (Collectives Acceleration Engines). Sa fonction est plutôt d’empiler les communications de tous les cœurs de calcul présents dans la puce pour les envoyer sur le réseau dans l’ordre le plus optimal possible. Selon Google, cet accélérateur CAE permettrait de diviser par 5 les temps de latence.
Sur Google Cloud, 1152 TPU 8i pourront faire partie du même cluster. Ici, il ne s’agit plus de faire travailler les puces ensemble sur un même calcul. Il s’agit plutôt de regrouper les différentes tâches d’une même application d’IA (agents, RAG…) sur des serveurs proches les uns des autres, de sorte à minimiser les trajets des données.
Pour approfondir sur Hardware IA (GPU, FPGA, etc.)
