Build on Trainium

Build on Trainium est un programme d’investissement de 110 millions USD axé sur la recherche en IA et l’enseignement universitaire afin de soutenir la prochaine génération d’innovation et de développement sur AWS Trainium. AWS Trainium est une puce systolique d’IA spécialement conçue pour développer des idées et des applications d’IA de pointe. Build on Trainium finance de nouvelles recherches en IA sur Trainium, en investissant dans des équipes universitaires de premier plan pour innover dans des domaines critiques, notamment de nouvelles architectures de modèles, des bibliothèques de machine learning, des optimisations, des systèmes distribués à grande échelle, etc. Cette initiative pluriannuelle jette les bases de l’avenir de l’IA en incitant la communauté universitaire à tirer parti de la communauté open source autour de Trainium, à y investir et à y contribuer. En combinant ces avantages avec le kit de développement logiciel (SDK) Neuron et le lancement récent de Neuron Kernel Interface (NKI), les clients de Trainium peuvent désormais innover à grande échelle dans le cloud.

Nous avons créé un cluster de recherche Trainium dédié comprenant jusqu’à 40 000 puces Trainium qui seront disponibles via des instances Amazon EC2 Trn1 connectées à un seul réseau non bloquant à l’échelle du péta-bit à l’aide d’ Amazon EC2 UltraClusters. Les équipes de recherche et les étudiants peuvent accéder à ces puces via des réservations de blocs de capacité autogérées à l’aide de blocs de capacité Amazon EC2 pour ML.

Nous organisons plusieurs séries d’appels à propositions (CFP) Amazon Research Awards (ARA) à l’intention de l’ensemble de la communauté des chercheurs. Les propositions sélectionnées recevront des crédits AWS Trainium et auront accès au cluster de recherche Trainium. Build on Trainium accueille les propositions de recherche qui tireront parti des bibliothèques et des cadres de ML open source bien connus, et contribueront à leur tour à l’open source pour améliorer les ressources destinées à la communauté des développeurs de ML.

Neuron Kernel Interface (NKI) est une nouvelle interface de programmation pour les puces AWS AI, Trainium et Inferentia. NKI fournit un accès direct aux primitives matérielles et aux instructions disponibles sur AWS Trainium et Inferentia, permettant aux chercheurs de créer et de régler des noyaux de calcul pour des performances optimales. Il s’agit d’un environnement de programmation basé sur Python qui adopte une syntaxe de type Triton et une sémantique au niveau des tuiles couramment utilisées. Les chercheurs peuvent utiliser le NKI pour améliorer les modèles de deep learning grâce à de nouvelles fonctionnalités, optimisations et innovations scientifiques. Consultez la page de documentation de NKI pour en savoir plus.

Au Device Realization Lab du MIT, nous utilisons AWS Trainium pour repousser les limites de la recherche sur l’IA médicale. Nos modèles 3D de segmentation par ultrasons et d’estimation de la vitesse du son s’entraînent plus rapidement et plus efficacement que jamais, réduisant de plus de moitié la durée des expériences tout en atteignant une précision de pointe. AWS Trainium nous a permis d’étendre nos recherches d’une manière qui n’était pas possible avec les systèmes GPU traditionnels. En entraînant nos réseaux neuronaux 3D entièrement convolutifs sur AWS Trainium (trn.32xlarge), nous avons atteint des performances de pointe avec un débit 50 % plus élevé et des coûts inférieurs à ceux des instances NVIDIA A100. À l’aide d’un cluster Trainium à 32 nœuds, nous avons mené efficacement plus de 180 expériences d’ablation, réduisant ainsi la durée totale de formation de plusieurs mois à plusieurs semaines et accélérant l’innovation en matière d’IA médicale au MIT. À l’avenir, nous prévoyons d’utiliser Trainium pour former des modèles d’agents d’IA capables de gérer et d’automatiser le flux de travail des ultrasons numériques, permettant ainsi aux cliniciens de gagner un temps considérable et de fournir de meilleurs soins aux patients. »

« Le groupe de recherche CMU Catalyst travaille à l’optimisation des systèmes de ML. Notre projet vise à faciliter l’optimisation des différents systèmes de ML. Trainium est unique en ce qu’il fournit à la fois un contrôle de bas niveau et une interface de programmation accessible via l’interface Neuron Kernel Interface (NKI).

avec le soutien d’AWS dans le cadre du programme Build on Trainium, notre chercheur a pu explorer des optimisations avancées sur un noyau critique : FlashAttention. Ce qui nous a le plus surpris, c’est la rapidité avec laquelle nous avons pu itérer : nous avons réalisé des améliorations significatives en seulement une semaine grâce à la documentation NKI, Neuron Profiler et d’architecture accessibles au public. La combinaison d’outils puissants et d’informations claires sur le matériel a permis à notre équipe d’accéder à une optimisation sophistiquée de bas niveau.

AWS Trainium et Neuron Kernel Interface (NKI) permettent aux chercheurs comme nous d’innover plus rapidement, en supprimant les obstacles qui ralentissent généralement les travaux d’optimisation spécifiques au matériel. »

 

« Grâce au programme Build on Trainium, son équipe a obtenu un accès complet à la nouvelle pile de compilateurs open source NKI d’AWS Neuron, y compris une visibilité directe sur le Trainium ISA et les API pour une planification et une allocation de mémoire précises. Ce niveau de visibilité et de contrôle permet à ses étudiants d’analyser plus facilement les opportunités d’optimisation et de découvrir plus efficacement des implémentations performantes. »

Christopher Fletcher, professeur agrégé d’informatique, Université de Californie à Berkeley
 

« L’accès à AWS Trainium et Inferentia a joué un rôle déterminant dans l’avancement de nos recherches et de notre formation sur les systèmes d’IA efficaces et à grande échelle. Nous utilisons ces plateformes pour la formation de plusieurs experts et les optimisations d’inférence, le prototypage de nouvelles techniques d’exécution et de planification qui améliorent la capacité de mise à l’échelle, l’efficacité et la portabilité sur les architectures d’accélérateurs émergentes. En tirant parti de la pile Neuron Developer, les chercheurs de l’UIUC développent de nouvelles techniques d’exécution et de planification qui améliorent l’efficacité et la portabilité des charges de travail d’IA. L’équipe est particulièrement impressionnée par l’ouverture de la pile Neuron Developer, qui rend ces plateformes précieuses pour la recherche en matière d’exécution et permet des innovations en matière de rareté, de hiérarchies de mémoire et d’efficacité des communications qui vont au-delà des architectures GPU traditionnelles. »

« En tirant parti d’AWS Trainium et du programme Build on Trainium, mes étudiants et moi avons pu accélérer considérablement nos simulations de circuits quantiques. Le projet a réuni un solide groupe d’étudiants qui ont construit en collaboration un simulateur de haute performance, permettant des expériences plus approfondies et un apprentissage pratique à une échelle qui n’était tout simplement pas possible auparavant. »

« Notre équipe de recherche de l’UTS explore l’intégration d’algorithmes de watermarketing en développant des noyaux Neuron NKI personnalisés. L’accès à la pile Neuron open source via le programme Build on Trainium a été transformateur. Cela nous donne une visibilité sans précédent sur l’architecture Trainium et la possibilité de travailler directement au niveau matériel. L’accès à Trainium a permis à notre équipe d’accélérer considérablement ses charges de travail en matière de watermarketing, en réduisant les cycles d’itération et en nous permettant d’explorer des modèles et des techniques plus complexes. Cette profondeur d’accès permet à nos chercheurs de mettre en place de nouvelles idées, d’expérimenter des optimisations de bas niveau et de repousser les limites de ce que les systèmes de watermarketing peuvent réaliser sur les accélérateurs d’IA modernes. »