Hugging Face annonce la prise en charge native de Trainium, la puce d’AWS, grâce à TorchNeuron, un nouveau backend pour PyTorch. Transformers, Diffusers et Kernels peuvent désormais s’exécuter directement sur cette puce, ce qui doit faciliter l’exécution des modèles sur Trainium et leur partage sur le Hub de Hugging Face. Selon la source, une version bêta publique sera disponible à partir du 20 octobre 2026.
- Hugging Face intègre Trainium nativement grâce à TorchNeuron, un backend pour PyTorch.
- Transformers, Diffusers et Kernels s’exécutent directement sur Trainium.
- Hugging Face recommande désormais ces bibliothèques plutôt qu’Optimum Neuron.
- Une version bêta publique arrive le 20 octobre 2026.
TorchNeuron est une extension de PyTorch, la bibliothèque de référence pour le développement de modèles d’apprentissage automatique. Elle ajoute à PyTorch un backend dédié à Trainium, avec le mode d’exécution « eager », torch.compile (un outil d’optimisation du code) et les API distribuées natives comme FSDP et TP. Les chercheurs et développeurs pourront l’installer avec la commande pip install torch-neuronx.
Un changement d’approche pour les bibliothèques Hugging Face
Jusqu’ici, l’utilisation de Trainium et d’Inferentia, deux puces conçues par AWS, l’une pour l’entraînement économique des modèles et l’autre pour l’inférence à grande échelle, passait par la bibliothèque optimum-neuron. Celle-ci reposait sur le paquet torch-neuronx, lui-même basé sur PyTorch XLA, ce qui imposait de raisonner « comme un compilateur » : tenseurs paresseux, formes statiques uniquement et recompilation à chaque rupture de graphe. Hugging Face indique qu’optimum-neuron servait à masquer cette complexité grâce à du code maintenu spécifiquement.
TorchNeuron modifie cette approche en s’intégrant directement à PyTorch, qui voit Neuron comme un périphérique natif. Trainium devient alors un accélérateur comme un autre pour Transformers, Diffusers, Accelerate, PEFT et TRL, avec les mêmes classes et scripts. Hugging Face recommande désormais aux utilisateurs de Neuron de passer par ces bibliothèques plutôt que par Optimum Neuron.
Pour la génération de texte, la prise en charge est large : selon Hugging Face, 89 % des 100 architectures de modèles de génération de texte les plus utilisées sont compatibles d’emblée, contre 46 % avec Optimum Neuron v0.4.5. Il suffit de passer l’argument device_map=neuron pour déplacer le modèle sur les appareils Neuron.
Entraînement, diffusion et noyaux de calcul
Pour l’entraînement, un script standard de fine-tuning avec TRL fonctionne sur une instance Trainium moyennant quelques réglages propres à Neuron. Hugging Face précise que l’exemple publié utilise Qwen/Qwen3-8B et le jeu de données HuggingFaceH4/ultrachat_200k, et qu’il a été exécuté sur 8 puces Trainium2, soit 32 NeuronCores, c’est-à-dire la moitié d’une instance trn2.48xlarge. Le point de sauvegarde produit est un checkpoint Transformers standard, utilisable sur d’autres matériels.
Plusieurs pipelines de Diffusers sont compatibles avec TorchNeuron, en mode « eager » comme compilé. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 et H3 supportent aussi le parallélisme de tenseurs, qui répartit chaque couche du modèle sur plusieurs cœurs. Hugging Face indique que d’autres pipelines seront progressivement compatibles et qu’un guide permet d’ajouter soi-même ce support.
Enfin, la bibliothèque Kernels accueillera des noyaux de calcul optimisés pour Neuron, écrits avec la Neuron Kernel Interface (NKI). Un noyau destiné aux transformers de vision, publié sur le Hub, permet selon la source une accélération « jusqu’à x2.4 » sur Trainium. Hugging Face remercie l’équipe d’AWS pour son travail sur TorchNeuron et sa réactivité face aux retours de la communauté.
