Amazon lance ses puces Trainium3 pour l’IA : 2,5 pétaflops et 144 Go de HBM3e

Amazon vient d’annoncer la disponibilité pour les clients de sa nouvelle puce Trainium3 pour l’entrainement et l’inférence des intelligences artificielles. La division cloud du géant américain continue sur sa lancée, avec une nouvelle génération par an.

Des UltraServers et UltraClusters avec Trainium3 disponibles

L’année dernière, durant la conférence re:Invent, c’était Trainium2 qui était disponible et la puce Trainium3 présentée. Encore un an avant (en 2023), Trainium2 était annoncée pour remplacer le SoC Trainium de première génération.

On rembobine cette fois avec les performances. Trainium2 était annoncée comme quatre fois plus rapide sur les entrainements, avec une capacité mémoire multipliée par trois. L'efficacité énergétique (performances/watt) était multipliée par deux.

Amazon promettait des instances Trn2 avec jusqu'à 100 000 puces Trainium2 via les UltraClusters, pour une puissance de calcul pouvant atteindre « jusqu'à 65 exaflops ». Des UltraServers plus raisonnables avec 64 Trainium2 interconnectés étaient aussi de la partie, pour une puissance maximale de 83,2 pétaflops.

Les UltraClusters Elastic Compute Cloud (EC2) « sont composés de milliers d'instances EC2 accélérées, colocalisées dans une zone de disponibilité AWS donnée et interconnectées via un réseau », explique Amazon. Les UltraServers sont d’au moins un ordre de grandeur plus petits, puisqu’il est question de « plusieurs instances EC2 ».

Trainium3 : x2 sur les performances, 50 % de mémoire en plus

Nous revoilà à Trainium3, dont les UltraServers équipés de cette génération proposent « une puissance de calcul jusqu'à 4,4 fois supérieure, une efficacité énergétique quatre fois plus grande et une bande passante mémoire presque quatre fois supérieure à celle des serveurs Trainium2 UltraServers ». Trainium3 est gravée avec une technologie 3 nm.

Un UltraServer peut héberger jusqu’à 144 puces Trainium3 pour une puissance de 362 pétaflops, avec une précision en virgule flottante de 8 bits (FP8). Avec un arrondi, on arrive bien au rapport x4,4 annoncé par Amazon par rapport aux 83,2 pétaflops de la génération précédente. Un UltraServer Trn3 complet dispose de 20,7 To de mémoire HBM3e et d’une bande passante agrégée de 706 To/s.

Avec une division par 144, cela nous donne 2,5 pétaflops de calculs en FP8 et 144 Go de mémoire par puce. Le même calcul donne respectivement 1,3 pétaflop et 96 Go pour la puce Trainium2. La bande passante mémoire totale de Trn2 est de 185 To/s (on a donc un gain de x3,8 avec Trainium3).

2 500 téraflops (ou 2,5 pétaflops) pour Trainium3 alors que NVIDIA revendique pour rappel 5 000 téraflops en FP8 pour un die Blackwell, mais le fabricant en intègre deux par puce, soit 10 000 téraflops. Pour rappel, la H100 de NVIDIA est annoncée pour 3 958 téraflops en FP8.

• NVIDIA annonce son GPU Blackwell (B200) pour l’IA, jusqu’à 20 000 TFLOPS (FP4)

UltraClusters : jusqu’à un million de puces Trainium3

L’entreprise affirme avoir amélioré sa partie réseau, notamment son NeuronSwitch-v1 qui permet de doubler la bande passante à l’intérieur des UltraServers. Pour les échanges entre les puces, une version améliorée de Neuron Fabric permet de réduire la latence à moins de 10 ms.

Les UltraClusters sont aussi améliorés et peuvent connecter entre eux des milliers d'UltraServers pour arriver à un total d’un million de puces Trainium, soit 10 fois plus que les 100 000 puces de la génération précédente.

Prudence sur les chiffres d'Amazon toutefois. La société se faisait récemment l’écho du projet Rainier : une « infrastructure collaborative avec près d’un demi-million de puces Trainium2 ». Trainium3 ne fait donc « que » deux fois mieux, pas 10. Les puces de Rainier sont réparties entre plusieurs centres de données aux États-Unis.

Selon Amazon, ses UltraClusters permettent « l'entraînement de modèles multimodaux avec des milliards de jetons et de l'inférence en temps réel pour des millions d'utilisateurs en simultané ».

Trainium4 : FP4 et au NVLink Fusion

Si vous avez lu attentivement le début de cet article, vous savez certainement de quoi parle la fin : Trainium4. Amazon annonce « six fois plus de performances de traitement (FP4), trois fois plus de performances en FP8 et quatre fois plus de bande passante mémoire ».

L’annonce sur FP4 est la conséquence directe du x3 en FP8. En divisant la précision par deux, on multiplie les performances par deux. Amazon rejoint donc NVIDIA qui était le premier à se lancer en FP4 avec Blackwell.

• Rapidité vs précision : deux experts nous expliquent les enjeux des GPU modernes sur les IA

Enfin, Amazon prévoit de rendre Trainium4 compatible avec NVLink Fusion, une technologie NVIDIA permettant de monter « une infrastructure d’IA semi-personnalisée », avec des puces du Caméléon et celles de partenaires. Au lancement, MediaTek, Marvell, Alchip Technologies, Astera Labs, Synopsys et Cadence s'étaient associés à l'annonce.