AWS lance une compétition Trainium Frontier pour co-concevoir modèles et kernels sur puces IA dédiées

Amazon Web Services a lancé une compétition baptisée AWS Trainium Frontier, destinée aux laboratoires universitaires et industriels souhaitant explorer l'entraînement de modèles de langage directement sur ses puces d'intelligence artificielle Trainium. Le principe consiste à entraîner un modèle de langage à partir de zéro, en partant d'une base fournie d'environ 50 millions de paramètres dérivée du projet nanochat, une architecture dense de type GPT utilisant la normalisation RMSNorm, les embeddings rotatifs et un module MLP à activation ReLU². Les participants peuvent modifier librement l'architecture, l'optimiseur, la boucle d'entraînement et même écrire des noyaux de calcul personnalisés via l'interface NKI (Neuron Kernel Interface). La première phase impose à chaque équipe une seule puce Trainium2 et un budget d'entraînement limité à trente minutes, un format assez court pour tester plusieurs dizaines d'hypothèses en une seule journée. Le score de cette phase repose sur un indicateur unique, le nombre de bits par octet en validation (val_bpb), mesuré à l'issue exacte des trente minutes. AWS met à disposition un support natif de PyTorch ainsi que des outils assistés par intelligence artificielle, incluant un accès à Amazon Bedrock, pour faciliter l'écriture de ces noyaux spécialisés.

L'enjeu dépasse la simple compétition technique: il s'agit de déterminer à quoi ressemble une architecture de modèle optimale lorsque les contraintes matérielles changent radicalement. Les puces Trainium disposent de davantage de mémoire embarquée (SBUF), d'un contrôle logiciel explicite sur les mouvements de données et de multiplications matricielles systoliques économes en énergie. Ce profil matériel modifie le rapport entre puissance de calcul et bande passante mémoire, si bien que des opérations habituellement limitées par la mémoire sur les accélérateurs classiques deviennent limitées par le calcul sur Trainium. Cela ouvre un espace de conception inédit, où il devient avantageux d'échanger davantage de calcul contre moins de trafic mémoire. Pour l'industrie, cela signifie que les schémas d'attention, les structures de MLP et les stratégies de parallélisme optimisés pour les GPU classiques ne sont peut-être plus les meilleurs choix sur ce type de matériel, ce qui pourrait redéfinir les pratiques de conception des futurs modèles.

Ce projet s'inscrit dans une tendance plus large où les architectures de modèles de langage se sont historiquement façonnées autour des contraintes des GPU dominants, qu'il s'agisse de la taille des warps, de la géométrie des tensor corés ou de la hiérarchie mémoire. En développant ses propres puces Trainium, AWS cherche à réduire sa dépendance à Nvidia tout en stimulant une recherche originale sur des architectures véritablement natives à son matériel. La simplicité de l'interface NKI, jugée accessible en un week-end aussi bien à un développeur qu'à un agent d'IA générant du code sous supervision humaine, doit permettre à un large éventail d'équipes de participer. Les meilleures propositions devront trouver l'équilibre entre capacité du modèle et vitesse d'entraînement dans un budget de temps fixe, révélant potentiellement des architectures radicalement différentes de celles conçues pour les GPU traditionnels.