Amazon a présenté lundi soir, un outil open source baptisé TorchNeuron, qui permet d’exécuter du code PyTorch sur ses puces IA Trainium sans le réécrire. Publié aujourd'hui, il s’attaque à l’un des atouts les plus solides du concurrent NVIDIA.

Avec TorchNeuron, Amazon veut faire de ses puces Trainium une vraie alternative pour les développeurs PyTorch. © Justlight / Shutterstock
Avec TorchNeuron, Amazon veut faire de ses puces Trainium une vraie alternative pour les développeurs PyTorch. © Justlight / Shutterstock

Avec TorchNeuron, le mastodonte de la vente en ligne Amazon veut rendre ses puces IA Trainium accessibles à tous les développeurs PyTorch, sans les obliger à réécrire leur code. Présenté le lundi 19 octobre, l’outil sera publié en open source dès aujourd'hui, à l’occasion de la PyTorch Conference North America. Plusieurs entreprises l’utilisent déjà, comme Hugging Face, des spécialistes de la voix ou en neurosciences, mais Amazon n’a encore communiqué aucun chiffre de performance face aux puces concurrentes.

Amazon s’attaque au vrai point fort de NVIDIA, et ce n’est pas une puce

Dans le monde de l'IA, presque tout le monde parle la même langue, et elle s'appelle PyTorch. Ce kit logiciel gratuit, qui sert à concevoir et à entraîner des modèles d'intelligence artificielle, serait utilisé par 85 % des chercheurs et développeurs du secteur, selon Amazon. Le problème, c'est qu'un programme écrit avec PyTorch ne s'exécute pas aussi facilement sur n'importe quelle puce. Pour changer de matériel, les équipes devaient souvent réécrire des pans entiers de leur code, puis tout retester pour s'assurer que leur IA se comportait toujours de la même façon. Un chantier de plusieurs mois, et un verrou qui a longtemps profité au leader du marché, NVIDIA, à la fois partenaire et concurrent d'Amazon Web Services, dont les puces sont depuis des années la cible naturelle de ces outils.

TorchNeuron fait sauter ce verrou pour Trainium, les puces conçues par Amazon et louées aux entreprises via son cloud AWS. Pour y parvenir, Amazon s'appuie sur une sorte de prise standard que PyTorch réserve justement aux nouveaux matériels. Il suffit d'y brancher sa puce pour qu'elle soit reconnue comme n'importe quelle autre. Conséquence directe, les outils dont les développeurs se servent au quotidien pour vérifier leur code, traquer les erreurs ou répartir le travail sur plusieurs machines fonctionnent sur Trainium sans la moindre retouche.

Amazon en a fait la démonstration en lançant le même programme d'entraînement sur une puce classique puis sur Trainium, en changeant uniquement le réglage qui indique quelle puce utiliser. Résultat, l'IA a appris les mêmes choses, au même rythme. Et pour éviter les mauvaises surprises, chaque nouvelle version de TorchNeuron est testée sur des centaines de milliers de modèles avant de toucher le moindre projet client.

Pour Amazon, le vrai coût de l'IA ne se résume pas au prix des puces. Dans beaucoup d'entreprises, et surtout dans les start-up, la dépense la plus lourde reste le temps des ingénieurs. Car chaque semaine passée à adapter un logiciel à un nouveau matériel est une semaine perdue pour le produit, ce qui peut peser sur la survie d'une petite équipe. Les grands groupes y trouvent un autre intérêt. Une IA vit deux phases : l'entraînement, quand elle apprend à partir de montagnes de données, et l'inférence, quand elle répond aux utilisateurs. Si les mêmes puces savent faire les deux, celles qui servent les clients aux heures de pointe peuvent entraîner de nouveaux modèles aux heures creuses, au lieu de tourner à vide. Enfin, si l'on peut passer d'une puce à l'autre sans pénalité, les fabricants devront séduire les développeurs plutôt que compter sur leur dépendance, avec à la clé des coûts plus bas et des projets qui avancent plus vite.

Le parcours d’une ligne de code écrite avec PyTorch. Elle est traduite par TorchNeuron, compilée par le logiciel AWS Neuron, puis exécutée sur la puce Trainium. © Amazon
Le parcours d’une ligne de code écrite avec PyTorch. Elle est traduite par TorchNeuron, compilée par le logiciel AWS Neuron, puis exécutée sur la puce Trainium. © Amazon

Amazon joue la carte de l’open source jusqu’au cœur de PyTorch

Avant même sa sortie officielle, TorchNeuron a trouvé preneur dans des secteurs comme la vidéo, la robotique, la voiture autonome, le traitement de la parole, la finance, les grands modèles de langage comme ChatGPT, ou encore les world models, ces IA qui apprennent à simuler le monde physique. La société Cartesia, qui conçoit des IA capables d'interagir avec leur environnement comme un humain, à commencer par la conversation, a fait tourner sur Trainium la chaîne logicielle qui produit les réponses de ses modèles, en ne retouchant son code qu'à la marge. La santé s'y met aussi. Le Lieber Institute for Brain Development, plus grand institut de recherche dédié aux traitements des troubles neuropsychiatriques comme la schizophrénie, mobilise l'IA pour dénicher de nouveaux médicaments. Il s'intéresse notamment aux simulations de dynamique moléculaire, qui reproduisent le mouvement des molécules atome par atome mais engloutissent des quantités colossales de calcul. Grâce à TorchNeuron, l'institut teste sur Trainium des IA capables de les accélérer, voire de les remplacer.

Amazon ne s'arrête pas à son propre outil et soigne tout l'écosystème qui l'entoure. Hugging Face, bibliothèque géante où les développeurs du monde entier partagent leurs modèles d'IA, et dont NVIDIA a pourtant annoncé le rachat début septembre, prend désormais en charge Trainium directement dans ses outils, sans exiger la moindre ligne de code propre à Amazon. TorchNeuron fonctionne aussi avec TorchTitan, l'outil officiel de PyTorch pour découper l'entraînement d'un modèle géant entre plusieurs machines, avec les mêmes méthodes que sur n'importe quelle autre puce.

Pour l'inférence, phase où l'intelligence artificielle répond aux utilisateurs, Amazon a apporté une contribution majeure à vLLM, un logiciel libre populaire qui accélère justement la production de ces réponses. Il a même publié l'ensemble de ses « kernels », ces petits programmes ultra-optimisés qui exécutent les calculs au plus près de la puce, et qui serviront de point de départ à ceux qui veulent écrire les leurs. De manière plus audacieuse, Amazon travaille avec la communauté PyTorch pour retirer du framework le code taillé pour des puces précises. Un PyTorch plus neutre profiterait à tous les développeurs, pas seulement aux clients de Trainium.

TorchNeuron sera publié sous licence Apache 2.0 ce mardi 20 octobre, lors de la PyTorch Conference North America, avec prise en charge des Trainium2 et Trainium3. Il prolonge une décennie d'investissements entamée avec le rachat d'Annapurna Labs en 2015, depuis lequel Amazon a conçu quatre familles de puces maison. Son activité silicium a d'ailleurs franchi les 25 milliards de dollars de revenus sur un an. La vitesse sera évidemment scrutée de près. Pour l'instant, c'est un secret de Polichinelle.