AWS Trainium est une famille de circuits intégrés spécifiques à une application (ASIC) développés par Amazon Web Services (AWS) pour accélérer les tâches d'apprentissage automatique et d'inférence dans le cloud. Annoncés en novembre 2020 et disponibles pour la première fois en 2022, les puces Trainium sont conçues pour offrir une alternative plus rentable et plus économe en énergie aux unités de traitement graphique (GPU) à usage général pour les charges de travail d'IA à grande échelle, en particulier celles impliquant des modèles d'apprentissage profond tels que les grands modèles de langage et les architectures transformeur. Les puces sont intégrées aux instances Elastic Compute Cloud (EC2) d'AWS, fournissant aux clients un matériel dédié optimisé pour les schémas de calcul spécifiques des opérations de réseaux neuronaux.
Le développement de Trainium s'inscrit dans une stratégie plus large d'AWS visant à diversifier ses offres matérielles d'IA et à réduire sa dépendance vis-à-vis des fournisseurs de puces externes comme Nvidia. En concevant son propre silicium, AWS vise à optimiser les performances par dollar et par watt pour les charges de travail d'IA dominantes exécutées sur sa plateforme. Trainium est complété par AWS Inferentia, une famille de puces distincte axée spécifiquement sur l'inférence, et ensemble, ils forment le cœur du portefeuille d'accélération d'IA personnalisée d'AWS. La famille Trainium a évolué avec des générations successives, y compris la deuxième génération Trainium2, annoncée en 2023 et déployée plus largement en 2024, et la troisième génération Trainium3, annoncée fin 2024.
Architecture matérielle
La puce Trainium repose sur une architecture spécialement conçue qui met l'accent sur la multiplication matricielle à haut débit et le mouvement des données, les opérations de base dans l'entraînement de réseaux neuronaux. Chaque puce Trainium contient plusieurs cœurs tensoriels capables d'effectuer des calculs en précision mixte, prenant en charge des types de données tels que FP32, FP16, BF16 et FP8. La puce Trainium de première génération, fabriquée selon un processus de 7 nanomètres par TSMC, offre jusqu'à 200 téraflops de performances pour les opérations FP16 et BF16. Elle comprend une grande mémoire sur puce et une interface à large bande passante vers la mémoire externe, permettant une gestion efficace des grands paramètres de modèles et des activations.
Trainium2, la deuxième génération, augmente considérablement les performances et la capacité mémoire. Fabriquée selon un processus de 5 nanomètres, chaque puce Trainium2 offre jusqu'à 400 téraflops de performances FP16/BF16 et dispose de 32 gigaoctets de mémoire à large bande passante (HBM). Un UltraServer Trainium2, qui est un système au niveau du rack, contient 64 puces Trainium2 interconnectées via une structure à haute vitesse, fournissant un total de 25,6 pétaflops de calcul et 6 téraoctets de mémoire. Cette conception permet l'entraînement de modèles avec des centaines de milliards de paramètres, tels que ceux utilisés dans les applications de IA générative.
Trainium3, annoncé en décembre 2024, devrait être fabriqué selon un processus de 3 nanomètres et devrait offrir environ deux fois les performances de Trainium2, avec un accent sur l'efficacité de l'entraînement et de l'inférence. AWS a déclaré que Trainium3 sera disponible en 2025, avec un accès anticipé fourni à certains clients.
Instances EC2 et disponibilité
AWS propose une capacité de calcul basée sur Trainium via son service EC2 sous les familles d'instances Trn1 et Trn2. Les instances Trn1 de première génération, lancées en octobre 2022, sont disponibles en deux tailles : trn1.2xlarge avec une puce Trainium et trn1.32xlarge avec 16 puces Trainium. La trn1.32xlarge offre jusqu'à 3,2 pétaflops de performances FP16/BF16 et 512 gigaoctets de mémoire, et prend en charge la mise en réseau à haute vitesse via l'Elastic Fabric Adapter (EFA) pour l'entraînement distribué sur plusieurs instances.
Les instances Trn2 de deuxième génération, lancées en décembre 2024, sont basées sur les puces Trainium2. L'instance trn2.48xlarge comprend 16 puces Trainium2, offrant 6,4 pétaflops de performances et 1,5 téraoctet de mémoire. AWS a également introduit l'UltraServer Trn2, qui combine quatre instances trn2.48xlarge en un seul système logique avec 64 puces Trainium2, permettant l'entraînement de modèles avec jusqu'à un billion de paramètres. Ces instances sont disponibles dans certaines régions AWS, notamment US East (Virginie du Nord), US West (Oregon) et Europe (Irlande), avec une expansion prévue.
Pile logicielle et intégration
Le matériel Trainium est pris en charge par le SDK AWS Neuron, un kit de développement logiciel comprenant un compilateur, un environnement d'exécution et des outils de profilage. Le SDK Neuron traduit les modèles écrits dans des frameworks populaires tels que PyTorch et TensorFlow en exécutables optimisés qui s'exécutent sur Trainium. Il prend en charge le parallélisme de modèles automatique, permettant de diviser les grands modèles sur plusieurs puces et instances avec une intervention minimale de l'utilisateur. AWS a également intégré Trainium à sa plateforme SageMaker, permettant aux clients de lancer des tâches d'entraînement sur des instances Trainium via la même interface gérée utilisée pour l'entraînement sur GPU.
En outre, AWS a développé une collaboration avec OpenAI et Anthropic pour optimiser leurs modèles respectifs pour Trainium. Anthropic, qui a un investissement significatif d'Amazon, a été un adoptant précoce clé, utilisant les UltraServers Trainium2 pour entraîner et déployer sa famille de modèles Claude. OpenAI a également annoncé en 2024 qu'il utiliserait Trainium2 pour entraîner certains de ses modèles de pointe, marquant une approbation notable des capacités de la puce personnalisée.
Performances et efficacité des coûts
AWS commercialise Trainium comme offrant jusqu'à 50 % de coût en moins par tâche d'entraînement par rapport aux instances EC2 comparables basées sur GPU, telles que celles alimentées par les GPU Nvidia A100 ou H100. Des benchmarks indépendants, tels que ceux de MLPerf, ont montré que Trainium2 atteint des performances compétitives sur les tâches d'entraînement standard, en particulier pour les modèles basés sur des transformeurs. Pour l'inférence, Trainium2 est conçu pour offrir un débit élevé avec une faible latence, ce qui le rend adapté aux applications en temps réel comme les chatbots et la génération de code.
Cependant, l'écosystème autour de Trainium est moins mature que celui de la plateforme CUDA de Nvidia. Certains frameworks et bibliothèques d'apprentissage automatique ont un support limité, et certaines fonctionnalités avancées, telles que l'attention sparse ou les noyaux personnalisés, peuvent nécessiter un effort d'ingénierie supplémentaire. AWS investit activement dans l'expansion des capacités du SDK Neuron et a publié une série de guides d'optimisation des performances pour aider les clients à atteindre une utilisation proche du pic.
Adoption et écosystème
Au-delà d'Anthropic et d'OpenAI, plusieurs autres organisations ont adopté Trainium pour leurs charges de travail d'IA. CoreWeave, un fournisseur de services cloud spécialisé dans l'infrastructure GPU, a annoncé en 2024 qu'il offrirait des instances basées sur Trainium à ses clients, élargissant la disponibilité de la puce au-delà du cloud d'AWS. Cerebras Systems, un concurrent dans le domaine des puces d'IA, a également établi un partenariat avec AWS pour offrir son propre matériel aux côtés de Trainium, bien que les deux produits ciblent des segments de marché différents.
Dans la communauté de recherche, Trainium a été utilisé pour entraîner de grands modèles dans des contextes académiques. Par exemple, des chercheurs de Stanford AI Lab et de Berkeley AI Research ont publié des études utilisant Trainium2 pour des expériences sur la mise à l'échelle et l'efficacité des modèles. AWS a également établi le programme de recherche Trainium, fournissant des crédits de calcul gratuits aux institutions académiques travaillant sur des projets d'IA open-source.
Paysage concurrentiel
Trainium concurrence directement d'autres accélérateurs d'IA personnalisés, y compris les unités de traitement tensoriel (TPU) de Google Cloud, les puces Maia de Microsoft Azure et les unités de traitement du langage de Groq. Alors que les TPU sont disponibles depuis plusieurs années et disposent d'une pile logicielle mature, l'avantage de Trainium réside dans son intégration étroite avec les services cloud étendus d'AWS et ses prix agressifs. L'émergence de Trainium3 devrait intensifier la concurrence, en particulier alors que la demande de calcul d'IA continue de dépasser l'offre.
L'investissement d'AWS dans le silicium personnalisé le positionne également contre AMD et Intel, qui offrent des options alternatives de GPU et d'accélérateurs, ainsi que contre les conceptions basées sur Arm Holdings qui gagnent du terrain dans les centres de données. Le succès de Trainium dépendra de la capacité d'AWS à maintenir un écosystème logiciel robuste et à démontrer des avantages clairs en termes de performances et de coûts à long terme.
Orientations futures
À l'avenir, AWS a indiqué que Trainium jouera un rôle central dans son projet Rainier, un supercalculateur massif construit en collaboration avec Anthropic. Le projet Rainier, annoncé en 2024, sera composé de centaines de milliers de puces Trainium2 et devrait être l'un des plus grands clusters d'entraînement d'IA au monde, avec une date opérationnelle prévue en 2025. Ce projet souligne l'engagement d'AWS envers le silicium personnalisé comme un atout stratégique pour capturer le marché croissant de l'infrastructure d'IA.
De plus, AWS explore l'utilisation de Trainium pour les déploiements en périphérie et sur site, bien qu'aucun produit spécifique n'ait été annoncé. L'entreprise investit également dans des technologies avancées de packaging et d'interconnexion pour étendre les systèmes Trainium au-delà du facteur de forme UltraServer actuel, permettant potentiellement l'entraînement de modèles avec des billions de paramètres.
Conclusion
AWS Trainium représente un pari significatif d'Amazon sur le silicium personnalisé pour l'IA. En offrant une puce spécialement conçue et étroitement intégrée à sa plateforme cloud, AWS vise à fournir aux clients une combinaison convaincante de performances, de coûts et d'évolutivité. Bien que des défis subsistent en matière de maturité logicielle et d'adoption de l'écosystème, l'évolution rapide de la famille Trainium et le soutien de grands laboratoires d'IA comme Anthropic et OpenAI suggèrent que Trainium sera un acteur majeur dans le paysage du matériel d'IA pour les années à venir.