La bêta de TorchNeuron apporte PyTorch en natif sur AWS Trainium

TorchNeuron, une nouvelle façon d'exécuter du code PyTorch sur les puces Trainium d'Amazon, est entré en bêta publique le 20 octobre 2026. C'est important, car la plupart des modèles ouverts populaires tournent désormais sur Trainium sans modification. Selon Hugging Face, 89 % des 100 architectures de modèles de génération de texte les plus utilisées fonctionnent telles quelles, contre 46 % avec les anciens outils.

Ce qu'est TorchNeuron

Trainium est la puce d'IA qu'Amazon conçoit et loue via AWS, son activité cloud. PyTorch est la bibliothèque la plus utilisée pour construire et entraîner des modèles d'IA. Jusqu'ici, le code PyTorch atteignait Trainium par une couche de traduction, et de nombreux modèles demandaient un travail supplémentaire pour fonctionner.

TorchNeuron est un plugin PyTorch qui ajoute Trainium comme périphérique natif. La documentation d'AWS indique qu'il « enregistre Trainium comme un périphérique PyTorch standard ». En pratique, un développeur déplace un modèle vers le périphérique nommé neuron, comme il le déplacerait vers un GPU avec cuda.

Le projet est open source sous licence Apache-2.0, selon sa page GitHub. AWS le présente comme une bêta publique dans la version Neuron 2.33, son kit logiciel pour ces puces.

Ce qui change par rapport à l'ancienne méthode

L'ancienne voie passait par Optimum Neuron, une bibliothèque de Hugging Face, qui reposait sur PyTorch XLA. XLA est une couche de compilation qui transforme un modèle en graphe figé avant son exécution. Cette approche rendait de nombreux modèles difficiles à prendre en charge, car chacun demandait sa propre adaptation.

TorchNeuron supprime cette couche. Les bibliothèques de Hugging Face tournent désormais directement sur Trainium : Transformers pour les modèles de langage, Diffusers pour les modèles d'images, ainsi qu'Accelerate, PEFT et TRL pour l'entraînement et le fine-tuning. Le tableau ci-dessous résume les fonctionnalités, d'après la documentation d'AWS et l'article de Hugging Face.

Fonctionnalité	Ce que cela signifie
Mode eager	Le code s'exécute ligne par ligne, comme en PyTorch classique
torch.compile	Compilation facultative pour gagner en vitesse, avec le backend neuron
Entraînement distribué	FSDP, DDP, DTensor, parallélisme de tenseurs et de pipeline
Kernels personnalisés	Code écrit à la main pour la puce via le langage NKI d'AWS
Puces	Trainium 2 et Trainium 3 ; FP8 et MXFP8 nécessitent Trainium 3

FSDP et DDP sont des méthodes standard de PyTorch pour répartir l'entraînement sur de nombreuses puces. FP8 est un format numérique plus petit qui économise de la mémoire et du temps.

Modèles et vitesse

L'article de Hugging Face montre un fine-tuning de Qwen3-8B avec TRL, réparti sur plusieurs puces grâce au parallélisme de tenseurs et à FSDP. Il fait aussi tourner des pipelines d'images comme FLUX.1-dev, FLUX.2 et Qwen-Image avec le parallélisme de tenseurs.

La documentation d'AWS liste des modèles mixture-of-experts, qui orientent chaque entrée vers quelques sous-réseaux spécialisés. DeepSeek, GPT-OSS et Qwen3 sont pris en charge. L'entraînement de ces modèles passe par un outil distinct, TorchTitan Neuron.

Le seul chiffre de vitesse a une portée étroite. Un kernel de vision transformer de l'équipe Neuron d'AWS offre « une accélération allant jusqu'à 2,4 fois sur Trainium », indique l'article de Hugging Face. Il ne précise pas par rapport à quoi elle est mesurée.

Comment l'installer

La page GitHub fournit une installation avec des versions figées, à lancer une fois le pilote et les paquets runtime de Neuron en place :

pip install torch==2.14.0 --index-url https://download.pytorch.org/whl/cpu
pip install "torch-neuronx==2.14.*" "neuronx-cc==2.*" "nki>=0.6" --extra-index-url https://pip.repos.neuron.amazonaws.com

Ces versions figées comptent. TorchNeuron est distribué sous le même nom de paquet, torch-neuronx, que l'ancien produit basé sur XLA. Les versions antérieures installent l'ancien, prévient le dépôt. Lancer python -m torch_neuronx.collect_env confirme que le périphérique est utilisable depuis PyTorch. Les images machine et les conteneurs prêts à l'emploi d'AWS l'intègrent déjà.

Ce que cela signifie pour les développeurs

TorchNeuron fait de Trainium non plus une cible à part, mais un périphérique PyTorch de plus. C'est l'étape qui rend les comparaisons de prix avec les GPU réalistes pour des équipes ordinaires.

• Testez d'abord votre propre modèle. Avec 89 % de couverture, environ une architecture populaire sur dix demande encore du travail. Lancez un court job d'entraînement ou d'inférence avant de planifier une migration.
• Figez les versions dans chaque environnement. Un torch-neuronx sans version figée dans un fichier requirements peut installer l'ancien produit XLA. Vérifiez la sortie de collect_env dans la CI.
• Vérifiez la génération de puce. Seules les instances Trainium 2 et 3 sont prises en charge. Si votre plan repose sur l'entraînement en FP8, il vous faut Trainium 3.
• Considérez-le comme une bêta. Aucune des trois sources ne le dit prêt pour la production. Gardez une solution de repli sur GPU pour tout ce qui a une échéance.
• Comparez le coût par job, pas par heure. Le chiffre de 2,4 fois concerne un seul kernel, sans base de comparaison indiquée. Mesurez vos propres tokens par dollar avant d'engager un budget.