Avec TorchNeuron, Amazon veut faciliter l’entraînement des modèles sur ses puces Trainium

PyTorch est une bibliothèque logicielle open source initialement développée par Meta (anciennement Facebook) en 2016. Elle permet de construire et d'entraîner des modèles d'intelligence artificielle. Depuis 2022, PyTorch est géré par la Linux Foundation, une organisation à but non lucratif qui soutient des projets logiciels libres. Cette bibliothèque est devenue un outil clé dans le domaine de l'IA, notamment grâce à sa flexibilité et sa facilité d'utilisation. Elle est comparée à TensorFlow, un autre framework populaire, mais se distingue par son approche plus intuitive pour les développeurs. PyTorch est particulièrement apprécié pour ses capacités de calcul automatique, qui simplifient la création de réseaux de neurones complexes.

CUDA est une technologie propriétaire développée par NVIDIA, leader mondial des cartes graphiques (GPU). Elle permet d'utiliser les GPU pour des calculs généraux (General-Purpose Computing on Graphics Processing Units, ou GPGPU), et pas seulement pour l'affichage graphique. Avant CUDA, les calculs complexes comme ceux nécessaires à l'entraînement des modèles d'IA étaient principalement effectués par les CPU (processeurs classiques), bien moins performants pour ce type de tâches. CUDA a révolutionné l'IA en rendant les GPU NVIDIA indispensables pour accélérer les calculs, créant ainsi une dépendance forte des développeurs à cette technologie. Cette domination a permis à NVIDIA de devenir un acteur central dans le domaine de l'IA.

Amazon Web Services (AWS) a développé ses propres puces dédiées à l'IA, appelées Trainium. La troisième génération de ces puces, lancée en décembre 2025, offre une puissance de calcul de 25 pétaflops (25 millions de milliards d'opérations par seconde) et dispose de 144 Go de mémoire HBM3E. Contrairement à NVIDIA qui se concentre uniquement sur les GPU, Trainium est conçu pour être utilisé à la fois pour l'entraînement des modèles d'IA et pour leur utilisation en production (inférence). Amazon n'est pas le seul à développer ses propres puces : Microsoft propose la puce Maia et Google utilise les TPU (Tensor Processing Units). Ces puces personnalisées visent à réduire la dépendance aux solutions de NVIDIA et à offrir des performances optimisées pour des tâches spécifiques.

TorchNeuron est un outil développé par AWS pour faciliter l'utilisation des puces Trainium avec la bibliothèque PyTorch. Il fonctionne comme une couche de traduction, similaire à CUDA pour NVIDIA. Son objectif est de permettre aux développeurs de remplacer facilement dans leur code les instructions destinées à CUDA par celles compatibles avec les puces Trainium. Par exemple, il suffit de remplacer le mot cuda par neuron dans le code pour utiliser Trainium. Cependant, cette simplicité est théorique : en pratique, certaines fonctions de PyTorch ne sont pas encore prises en charge par TorchNeuron, et des tests officiels de PyTorch peuvent être ignorés faute de traduction complète. Malgré ces limites, AWS met en avant des gains significatifs, comme une augmentation de 70 % de la quantité d'audio traitée par dollar dépensé pour une entreprise comme Deepgram, bien que sa mise en œuvre ait nécessité plusieurs mois de travail.

AWS présente TorchNeuron comme un moyen de réduire les coûts de changement de puce pour les développeurs, leur offrant ainsi plus de liberté pour comparer différentes options matérielles. Colin Brace, responsable chez Annapurna Labs (la filiale d'AWS qui conçoit les puces Trainium), souligne que l'infrastructure et l'écosystème d'AWS, ainsi que la possibilité d'utiliser la même puce pour l'entraînement et l'inférence, sont des atouts majeurs. Cette flexibilité permet de réaffecter la capacité inutilisée des puces, optimisant ainsi les ressources. Par ailleurs, AWS travaille directement avec la communauté PyTorch pour rendre cette bibliothèque plus indépendante du matériel, ce qui pourrait bénéficier à tous les développeurs, pas seulement aux utilisateurs de Trainium. Cette approche vise clairement à réduire la dépendance à NVIDIA, dont l'écosystème est souvent critiqué pour sa rigidité.