AWS Trainium gagne un backend natif dans PyTorch

Hugging Face prépare TorchNeuron, un backend natif pour PyTorch sur AWS Trainium, avec une disponibilité en bêta publique annoncée le 20 octobre 2026. L’enjeu est simple : rendre l’usage de Trainium plus proche de celui des GPU classiques, sans passer par une couche spécifique comme Optimum Neuron. Les bibliothèques Transformers, Diffusers, Accelerate, PEFT et TRL sont concernées, ce qui change la donne pour l’entraînement comme pour l’inférence à grande échelle.

Trainium cesse d’être un cas à part

Jusqu’ici, l’expérience Hugging Face sur Trainium et Inferentia passait par optimum-neuron, une couche bâtie sur torch-neuronx et PyTorch XLA. Cette approche fonctionnait, mais elle imposait des contraintes de compilation et de formes statiques, avec des recompilations fréquentes au moindre “graph break”. TorchNeuron inverse la logique : Trainium devient un périphérique PyTorch natif, avec exécution eager, torch.compile et primitives distribuées standard comme DTensor, FSDP2 et FSDP.

Autrement dit, les développeurs n’ont plus à adapter leurs scripts à une logique de compilateur. Les mêmes classes et les mêmes appels peuvent désormais fonctionner sur plusieurs accélérateurs. Hugging Face recommande même de passer directement par ses bibliothèques plutôt que par Optimum Neuron. Les vieux contorsionnistes du code peuvent souffler un peu.

Transformers et Diffusers montent en charge

Dans Transformers, la génération de texte est désormais prise en charge nativement pour un large ensemble de modèles. Selon Hugging Face, 89 % des 100 architectures de génération de texte les plus utilisées sont supportées “out of the box” avec TorchNeuron, contre 46 % avec Optimum Neuron v0.4.5. Côté usage, il suffit d’ajouter device_map=neuron pour déplacer le modèle vers les Neuron devices.

Pour l’entraînement, la logique reste la même : un script TRL standard peut tourner sur Trainium avec seulement quelques réglages liés à Neuron. Hugging Face précise que le modèle est sharded à l’arrivée des poids grâce à une DistributedConfig combinant tensor parallelism et FSDP, sans wrapper supplémentaire du Trainer. Le support de torch.compile reste actif, tandis que le chargement asynchrone des poids doit être désactivé pour l’instant avec HF_DEACTIVATE_ASYNC_LOAD=1.

La partie image suit le même mouvement. Plusieurs pipelines de Diffusers sont compatibles en mode eager et compilé, dont FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 et H3. Hugging Face met aussi en avant le tensor parallelism, qui permet de répartir un même modèle sur plusieurs NeuronCore pour mieux exploiter la machine.

Ce que cela change pour les équipes qui entraînent vraiment des modèles

Le point le plus concret n’est pas le vernis technique, mais la réduction de friction. Un pipeline qui utilisait déjà TRL sur GPU peut, dans de nombreux cas, être transposé sur Trainium sans réécriture profonde. La contrainte reste réelle — formes statiques, padding à multiples de 128 tokens avec TRL ≥ 1.16, cache statique recommandé pour limiter les recompilations — mais l’architecture devient plus lisible et plus proche des habitudes PyTorch.

Hugging Face dit avoir exécuté un entraînement sur 8 puces Trainium2, soit 32 NeuronCores, l’équivalent de la moitié d’une instance trn2.48xlarge. Le résultat est ensuite sauvegardé dans un checkpoint Transformers standard, que l’on peut pousser sur le Hugging Face Hub et réutiliser sur d’autres matériels. C’est là que le sujet compte vraiment : réduire le verrouillage matériel sans perdre les gains de coût promis par Trainium.

Les kernels Neuron entrent eux aussi dans la boucle

Au-delà des modèles eux-mêmes, Hugging Face pousse aussi les kernels Neuron, écrits via le Neuron Kernel Interface (NKI). Ces kernels servent à optimiser des opérations précises, avec un contrôle plus bas niveau que les bibliothèques habituelles. L’idée est de publier davantage de composants réutilisables sur le Hub, y compris pour des opérations custom.

Hugging Face cite un exemple de gain pouvant atteindre x2,4 sur des Vision Transformers avec ce type de kernel. C’est une promesse utile, à condition de rester prudents sur le périmètre exact de mesure [à vérifier]. En pratique, le message est clair : Trainium ne doit plus ressembler à une île isolée, mais à un accélérateur de plus dans la boîte à outils PyTorch. AWS, de son côté, gagne un argument de poids face aux GPU dominants du marché.