Puces IA : Google élargit sa gamme de TPU, pour les entreprises

Des accélérateurs pour des usages spécifiques

Les TPU 8t et 8i annoncés en avril dernier étaient, comme leurs lettres l’indiquent, respectivement conçus pour l’entraînement (training) et l’inférence. Le premier a été pensé pour les traitements massivement parallèles en cluster, tandis que le second est optimisé pour les grandes fenêtres de contexte.

Le TPU 8t a moins de mémoire (128 Mo de cache SRAM et 216 Go de mémoire HBM), mais davantage de performances (12,6 pétaflops en FP4), et il peut être déployé en très grand nombre dans un cluster (9600 puces, avec une seule tâche qui peut être répartie sur 8640 puces). Il lit/écrit ses données en mémoire mémoire HBM à la vitesse de 6,5 To/s et communique avec d’autres TPU 8t à la vitesse de 2,4 To/s, à raison de liens directs en 400 Gbit/s entre deux TPI 8t. 

Le TPU 8i a davantage de mémoire (384 Mo de cache SRAM et 288 Go de mémoire HBM), mais moins de performances (10,1 pétaflops en FP4) et on ne peut en déployer que 1152 exemplaires dans un cluster, avec un maximum de 1024 puces dédiées à une seule tâche. Il lit/écrit ses données en mémoire HBM à la vitesse de 8,6 To/s et communique avec d’autres TPU 8i à la vitesse, là aussi, de 2,4 To/s, mais à raison de liens réseau RoCE en 200 Gbit/s entre deux TPU 8i.

Le nouveau TPU 8tx qui, contrairement à ce que sa première lettre indique, sert donc à l’inférence, est globalement la même chose que le TPU 8t. À deux différences près : les cœurs ARM Axion de Google sont ici remplacés par des cœurs x86 Zen d’AMD, pour apporter davantage de puissance de calcul afin de travailler avec des données chiffrées. Comme pour le TPU  8t, Google ne précise pas le nombre de ces cœurs. Il ne dit pas non plus de quelle génération de cœurs Zen il s’agit.

Ce chiffrement au plus bas niveau relève d’une discipline appelée informatique confidentielle. La technique retenue par Intel et AMD consiste à charger en RAM une machine virtuelle (VM) déjà chiffrée et à n’en déchiffrer des morceaux que le temps de leur traitement dans le processeur. Cela fonctionne aussi pour les applications en containers, car le système les charge dans une mini VM. Cette approche empêche le propriétaire des serveurs sur lesquels s’exécute un traitement d’espionner les données de ses clients en faisant une copie brute de la RAM de ses machines. Sans cette technique, les données ne sont chiffrées que sur les disques et sont déchiffrées dès leur chargement en mémoire.

La seconde différence est que les 2,4 To/s de bande passante vers l’extérieur sont, cette fois-ci, portés par des liens réseau en 200 Gbit/s. Comme sur le TPU 8i. Mais avec la particularité que ces liens sont à présent chiffrés en temps réel, pour éviter un espionnage des données par des sondes réseau lorsque deux TPU 8tx s’échangent des données pour se répartir le travail.

En ce qui le concerne, le nouveau TPU 8i Max est une version étendue du TPI 8i. Comme ses performances (20,1 pétaflops en FP4, bande passante HBM de 17,2 To/s), sa mémoire cache SRAM (768 Mo) et la bande passante de ses liens réseau (400 Gbit/s en RoCE) sont doublées, la puce possède très certainement deux fois plus de circuits de cœurs et de contrôleurs que le TPU 8i. En revanche, sa mémoire HBM n’a été étendue qu’à 384 Go (1,33 fois la capacité d’origine). Et, là aussi, on note que les cœurs ARM Axion ont été remplacés par des cœurs x86 Zen d’AMD.

Selon Google, l’intérêt du TPU 8i Max n’est pas tant d’exécuter simultanément deux fois plus de tâches d’inférence. Il s’agit plutôt de savoir traiter des données plus lourdes dans la fenêtre de contexte. Le fournisseur donne l’exemple des fichiers vidéo.

Des logiciels pour tirer pleinement parti de son propre cluster de TPU

Suite à l’adoption des TPU par OCI, des analystes prédisaient qu’il serait complexe d’exécuter dessus des codes d’IA compilés avec des outils qui ne savent les optimiser que pour les GPU de Nvidia ou d’AMD. Google contourne le problème en mettant désormais en Open source le code de TorchTPU, une bibliothèque qui permet l'exécution native sur les TPU des codes PyTorch et des modèles d’IA distribués sur Hugging Face. Il existe aussi des versions natives pour TPU des modules vLLM et SGLang qui servent à coder des outils d’inférence et d’interaction avec les LLM.

La disponibilité de ces logiciels témoigne une fois de plus de la volonté de Google de laisser les entreprises déployer elles-mêmes des TPU en leur donnant les outils pour les exploiter au maximum. Au sein de Google Cloud, seul l’hyperscaler était concerné par cette problématique, les clients ne pouvant exécuter sur les TPU en ligne que des applicatifs prêts à l’emploi.

Parmi les autres logiciels accompagnant les TPU de Google, un module d’exécution Agent Compute Engine suspend automatiquement les sessions d'agents ou d’utilisateurs inactifs, en déplaçant leur mémoire de travail hors de la puce en quelques millisecondes. Ce module repose sur une bibliothèque TPU Sync qui unifie sur le même espace d’adressage la mémoire HBM interne à la puce, la mémoire DRAM du serveur hôte et les unités de stockage NVME auquel ce serveur a accès. Cette architecture réduirait la latence habituelle du transfert des KV Caches de plus de 90 %.

Manifestement, cette technologie est compatible avec toutes les baies de stockage, au contraire de STX, l’approche similaire de Nvidia.

Enfin, concernant le fonctionnement en cluster, il est possible de subdiviser à la demande d’un pool de TPU en tranches (slices) grâce au moteur GKE dynamic Slicing qui prend la forme d’un container Kubernetes. Selon Amin Vahdat, ce dispositif éviterait aux utilisateurs de TPU de devoir reconfigurer à la main des silos de puces à chaque fois qu’ils souhaitent redéployer une nouvelle panoplie d’IA avec des exigences différentes. GKE Dynamic Slicing éviterait surtout de mobiliser sur une tâche des TPU qui ne sont pas sollicités et ainsi pouvoir exploiter leur puissance sur d’autres tâches.