Amazon vient d’ouvrir la bêta publique de TorchNeuron, un backend natif permettant d’utiliser PyTorch sur ses accélérateurs Trainium2 et Trainium3. Dans le cas le plus simple, déplacer un modèle vers la puce demande seulement d’ajouter .to('neuron'). C’est une avancée importante, mais cela ne veut pas dire que n’importe quel projet fonctionnera immédiatement sans vérification.
TorchNeuron sert à quoi ?
Trainium est la famille de puces conçue par Amazon pour entraîner et exécuter des modèles d’intelligence artificielle dans AWS. Le problème, c’est que l’écosystème logiciel reste aussi important que le matériel.
Une grande partie de la recherche et des applications d’IA repose aujourd’hui sur PyTorch. NVIDIA a construit une énorme avance grâce à CUDA et aux bibliothèques qui gravitent autour. Amazon doit donc permettre aux développeurs de conserver leurs outils habituels s’il veut les convaincre d’utiliser Trainium.
C’est exactement le rôle de TorchNeuron. D’après la présentation publiée par PyTorch, Trainium apparaît désormais comme un périphérique PyTorch standard nommé neuron.
Un modèle peut donc être déplacé vers l’accélérateur avec une instruction très familière :
model = MyModel().to("neuron")Le backend prend en charge le mode d’exécution immédiate, appelé eager, ainsi que torch.compile, les outils distribués et le profiler de PyTorch.
Est-ce vraiment compatible sans réécrire son code ?
Pour du code moderne construit avec les API standard de PyTorch, la migration peut être très légère. Les projets utilisant torch.accelerator, l’interface recommandée pour rester indépendant du matériel, peuvent même fonctionner sans modification du code principal.
Mais il y a forcément des limites.
Un modèle rempli d’opérations personnalisées, d’extensions CUDA ou d’optimisations propres aux GPU NVIDIA ne va pas devenir compatible comme par magie. Il faudra vérifier la prise en charge de chaque composant, mesurer les performances et parfois écrire des noyaux spécifiques avec la Neuron Kernel Interface.
Même la compilation demande de désigner le backend Neuron. La promesse « sans réécriture » est donc réaliste pour les projets qui respectent déjà les API communes, pas pour tous les programmes PyTorch existants.
Hugging Face fait partie de l’équation
TorchNeuron s’intègre également aux bibliothèques les plus utilisées de Hugging Face. Transformers, Diffusers, PEFT et TRL peuvent exploiter le nouveau backend, ce qui ouvre la porte aux modèles de langage, à la génération d’images et aux méthodes d’ajustement plus légères.
Hugging Face explique dans sa documentation consacrée à TorchNeuron que Trainium peut maintenant être détecté comme les autres accélérateurs. Pour Amazon, c’est essentiel : personne n’a envie de réécrire toute sa chaîne d’entraînement uniquement pour tester une nouvelle puce.
Le projet est aussi open source sous licence Apache 2.0. Le dépôt officiel TorchNeuron est disponible publiquement depuis le 20 octobre, tandis que l’installation passe par le dépôt de paquets AWS Neuron.
Amazon attaque le vrai avantage de NVIDIA
La bataille ne se joue pas uniquement sur le nombre de cœurs ou la puissance théorique. NVIDIA domine parce que CUDA est installé dans les habitudes des chercheurs, des développeurs et des entreprises.
Amazon essaie donc de supprimer le coût du changement. Cette logique rappelle les efforts de Microsoft avec sa puce Maia 300 pour l’inférence dans Azure. Les géants du cloud développent leurs propres accélérateurs, mais ils savent qu’un bon matériel sans outils familiers restera difficile à adopter.
TorchNeuron ne garantit pas que Trainium deviendra automatiquement plus rapide ou plus rentable qu’un GPU NVIDIA pour chaque modèle. Par contre, il élimine un obstacle très concret : devoir apprendre une nouvelle manière de programmer avant même de commencer les tests.
C’est probablement la partie la plus importante de cette annonce. Amazon ne demande plus aux développeurs de quitter PyTorch pour essayer Trainium. Il essaie de faire entrer Trainium directement dans PyTorch.