Google profite de sa conférence Next ’23 pour présenter ses nouveaux TPU de 5e génération : les v5e. Ils ne misent pas sur la performance brute (ils sont d’ailleurs en dessous des v4), mais sur la rentabilité avec une forte amélioration du rapport performances/prix. On a aussi les TPU v4e et v5 sur Google Kubernetes Engine et Multislice.

C’est quoi un TPU ?

Avant d’entrer dans le vif du sujet, un rappel important sur les TPU (Tensor Processing Unit) de Google. Ce sont « des circuits intégrés spécifiques aux applications (ASIC) conçus par Google pour accélérer les charges de travail de machine learning ». On peut les utiliser via des frameworks tels que TensorFlow, Pytorch et JAX.

Dans chaque TPU, on retrouve un ou plusieurs TensorCore (cela évolue en fonction des générations et des versions), eux-mêmes constitués d'une ou de plusieurs unités matricielles, d'une unité vectorielle et d'une unité scalaire. Plusieurs puces TPU peuvent aussi être interconnectées entre elles pour former un pod, que l’on peut de nouveau assembler. De plus amples détails sont disponibles sur cette page 

Les premiers TPU ont été annoncés par Google en 2016, puis nous avons eu les v2 (2017), v3 (2018) et v4 (2021). Cette année, le géant du Net présente le premier représentant de sa 5e génération de Tensor Processing Unit : le v5e. Selon Google, il s’agit du « compromis idéal entre performances et rentabilité ». 

Voici les TPU v5e, avec un meilleur ratio perfs/dollar

À défaut de détails techniques, on a droit à quelques promesses dans le communiqué : « par rapport aux TPU v4, les v5e offrent des performances sur l’entrainement des modèles jusqu'à deux fois supérieures par dollar et des performances d'inférence jusqu'à 2,5 fois supérieures par dollar pour les LLM et les modèles d’IA générative ».

On avait droit à un discours bien différent lors du passage de la v3 à la v4 : « Les TPU v4 surpassent les v3 de 2,1 fois en moyenne par puce et améliore le ratio performances/watt de 2,7 fois ». Il était alors question de performances brutes et d’efficacité énergétique. Avec les v5e, Google met en avant la rentabilité. Nous allons voir que la nuance est importante. 

TPU v5e du côté technique : un TensorCore, quatre MXU

Alors que les TPU v3 et v4 contiennent chacun deux TensorCore, le v5e n’en a qu’un seul. Chaque TensorCore de la puce (que ce soit en v3, v4 ou v5e) comporte quatre unités de multiplication matricielle (MXU), une unité vectorielle et une unité scalaire.

Afin de mieux comprendre la différence entre les v4 et v5e, voici un diagramme des deux puces :

Diagramme du TPU v5e à gauche et du v4 à droite. On voit bien la différence de nombre de TensorCore

Dans le but de comparer plus rapidement les caractéristiques techniques des différentes générations, nous avons regroupé dans un tableau celles que nous connaissons pour le moment sur les v5e, aux côtés des v3 et v4. 

On peut d’ores et déjà voir que la quantité de mémoire HBM2 est divisée par deux (mais il y a aussi deux fois moins de TensorCore dans les v5e) avec une bande passante inférieure à celle des TPU v3 et v4. La signification du « e » n’est pas précisée, mais cela pourrait être pour « efficience ».

Des péta/tera/exa opérations par seconde

Quoi qu’il en soit, jusqu’à 256 puces TPU v5e peuvent être interconnectées pour former un pod, avec une bande passante totale de 400 Tb/s, pour des performances de 100 PetaOP (OP pour opérations par seconde) en INT8 (des entiers sur 8 bits), selon le communiqué de Google.

Sur cette page, le géant du Net revendique 197 TeraFLOP (FLOP pour opérations en virgule flottante par seconde) en Bfloat16 et 393 TeraOP en INT8. Si on multiplie 393 TerOP par 256 puces TPU v5e, on arrive bien aux 100 PetaOP annoncés pour le pod (vous suivez ?).

• Les nombres en informatique : entiers, virgule flottante, simple et double précision (FP32/64), Tensor Float (TF32)…

Ce n’est pas la même chanson avec les TPU v4. Ils sont annoncés par Google à 275 TFLOP en Bfloat16 (bien au-dessus des v5e donc)… mais aussi à 275 TFLOP en INT8, et donc en dessous des v5e dans ce cas.

Avec 4 096 puces et donc 1 126 PetaFLOP de puissance totale (Bfloat16 ou INT8), un pod de TPU v4 dépasse largement un pod de TPU v5 qui arrive « seulement » à 50 PetaFLOP en Bfloat16 ou 100 PetaFLOP en INT8. Dans le cas des entiers, les performances sont divisées par 10, mais il y a 16 fois moins de puces.

Cette puce V5e laisse donc de la place à Google pour lancer un TPU v5 « complet »… avec (au moins) deux TensorCore comme les v4 et v3 ? Selon Dylan Patel de SemiAnalysis, les TPU v5 (alias Viperfish) pourraient également être regroupés par paquet de 16 384, si le rythme actuel est maintenu : 256 sur les v2, 1 024 sur les v3, 4 096 sur les v4, soit x4 à chaque génération.

Des travaux sur TPU v5 avaient été annoncés en juin 2021 par Anna Goldie, qui était alors chez Google Brain. De telles puces peuvent très bien être utilisées par l’entreprise, sans avoir été annoncées pour le moment. 

Google Kubernetes Engine sur TPU v4 et v5e et Multislice

À l’occasion de l’annonce des TPU v5e, l’entreprise en profite pour faire une annonce autour de Google Kubernetes Engine (GKE) : « Les clients peuvent maintenant améliorer la productivité lors de développement de l'IA en exploitant GKE pour gérer l'orchestration des charges de travail à grande échelle sur TPU v5e et v4 ».

Google présente aussi sa technologie Multislice. Elle était déjà utilisée en interne pour les grands modèles de langage (LLM) et débarque désormais en preview pour tous les clients. « Jusqu'à présent, les entrainements utilisant des TPU étaient limitées à une seule tranche de puces TPU, limitant la taille des plus gros entrainements à une tranche de 3 072 puces au maximum avec les TPU v4 ».

Dorénavant, il est possible d’exploiter « jusqu'à des dizaines de milliers de puces via inter-chip interconnect (ICI) au sein d'un seul pod, ou bien sur plusieurs pods via le datacenter network (DCN) ». Multislice n’est compatible qu’avec les TPU v4, v5e et le framwork JAX précise la documentation.