TorchNeuron simplifie l’usage de Trainium pour PyTorch et les modèles HF

AWS a présenté TorchNeuron, un nouveau backend natif pour PyTorch sur Trainium, ses puces d’IA destinées à l’entraînement comme à l’inférence. L’outil, annoncé le 19 octobre 2026 et prévu en bêta publique le 20 octobre, vise à rendre l’exécution de modèles plus simple pour les développeurs Hugging Face, sans réécriture lourde du code. Dans le même mouvement, Amazon attaque l’avantage logiciel de Nvidia, encore très lié à l’usage de PyTorch.

Trainium veut sortir du laboratoire

Jusqu’ici, l’expérience Hugging Face sur Trainium et Inferentia passait surtout par optimum-neuron, une couche qui masquait la complexité de PyTorch XLA mais imposait une logique de compilation plus contraignante. TorchNeuron change la méthode : le backend expose Neuron comme un périphérique natif, avec exécution eager, torch.compile et primitives distribuées PyTorch comme FSDP et DTensor. Pour les équipes, cela ressemble moins à un bricolage spécialisé qu’à une intégration standard. C’est un détail technique, mais en IA ce genre de détail décide souvent qui gagne du temps et qui le perd.

Hugging Face recommande désormais d’utiliser directement Transformers, Diffusers, Accelerate, PEFT et TRL avec Trainium. L’intérêt est simple : les mêmes scripts peuvent s’exécuter sur Trainium comme sur d’autres accélérateurs, avec seulement quelques réglages spécifiques à Neuron. AWS défend aussi une logique économique plus large, en présentant Trainium comme une plateforme unique pour l’entraînement et l’inférence, intégrée à son cloud.

Des gains de compatibilité, et moins de code spécial à maintenir

Dans Transformers, l’inférence textuelle est désormais prise en charge nativement pour un large éventail de modèles. Hugging Face affirme que 89 % des 100 architectures de génération de texte les plus utilisées sont prises en charge avec TorchNeuron, contre 46 % avec Optimum Neuron v0.4.5. L’activation tient à une simple option, device_map=neuron, selon la documentation publiée avec l’annonce. Là encore, l’enjeu n’est pas seulement la performance brute, mais la baisse du coût d’adaptation pour les équipes produit.

Du côté de l’entraînement, les scripts TRL peuvent fonctionner sur Trainium sans trainer dédié ni fork spécifique du modèle. L’exemple fourni par Hugging Face s’appuie sur SFTTrainer, avec parallélisme géré nativement par DistributedConfig. Les auteurs précisent toutefois quelques contraintes : shapes statiques, désactivation de certaines formes de dynamic shapes dans torch.compile, et chargement asynchrone des poids désactivé via HF_DEACTIVATE_ASYNC_LOAD=1. Bref, ça avance, mais les puces aiment encore un peu les contours nets.

Les bibliothèques HF suivent, puis les usages

Le support couvre aussi Diffusers pour l’image générative. Plusieurs pipelines, dont FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 et H3, sont compatibles avec le parallélisme tensoriel. Hugging Face montre un exemple avec une image générée sur Trainium en mode distribué, en laissant l’encodeur texte sur CPU et en compilant les blocs répétés côté Neuron. L’objectif est clair : faire de Trainium une cible de plus, et non une exception qu’il faut traiter à part.

AWS dit aussi vouloir contribuer du code directement au projet PyTorch pour réduire les dépendances implicites à Nvidia. De son point de vue, l’enjeu dépasse Trainium : il s’agit de rendre PyTorch plus matériel-agnostique. TechCrunch rapporte qu’Amazon voit là un moyen de challenger plus frontalement le couple Nvidia-CUDA, qui reste la référence de fait pour beaucoup d’équipes IA. C’est une bataille logicielle autant que matérielle, et elle compte car le choix d’un framework détermine ensuite celui du cloud, du processeur et parfois du budget.