TorchNeuron simplifie les modèles HF sur AWS Trainium

Hugging Face a détaillé TorchNeuron, une nouvelle intégration PyTorch pour les puces AWS Trainium, conçue pour simplifier l’exécution de modèles sur cet accélérateur dédié à l’entraînement IA. Ce module, prévu en bêta publique à partir du 20 octobre 2026 et installable avec pip install torch-neuronx, expose Neuron comme une cible native de PyTorch. Il prend en charge l’exécution immédiate, la compilation via torch.compile ainsi que les API distribuées de PyTorch, dont DTensor, FSDP2 et le parallélisme tensoriel. Auparavant, l’utilisation des modèles Hugging Face sur Trainium reposait principalement sur Optimum Neuron et sur PyTorch XLA. Cette approche nécessitait souvent de fonctionner avec des tenseurs paresseux, des formes fixes et des recompilations lors de certaines ruptures de graphe. Optimum Neuron masquait cette complexité grâce à des composants spécifiques, mais cela compliquait l’adoption par les développeurs habitués aux outils PyTorch standards. TorchNeuron réduit fortement ces frictions. Les bibliothèques Hugging Face telles que Transformers, Diffusers, Accelerate, PEFT et TRL peuvent désormais être utilisées directement sur Trainium, avec les mêmes classes et scripts que sur d’autres accélérateurs. Hugging Face recommande aux utilisateurs Neuron de recourir directement à ces bibliothèques plutôt qu’à Optimum Neuron. Pour l’inférence de génération de texte, TorchNeuron prend en charge 89% des 100 architectures les plus utilisées, contre 46% avec Optimum Neuron v0.4.5. Le modèle peut être placé sur les périphériques Neuron via un simple argument de placement, tandis que le parallélisme tensoriel permet d’exploiter plusieurs NeuronCores. L’utilisation d’un cache statique est également recommandée pour limiter les recompilations pendant la génération. L’entraînement bénéficie de la même simplification. Un script standard de réglage fin avec TRL peut fonctionner sur Trainium en transmettant une configuration distribuée pour la répartition tensorielle et FSDP. Les poids sont découpés au chargement du modèle, sans fork spécifique ni drapeau dédié. Le checkpoint final reste un format Transformers standard, partageable sur le hub Hugging Face et rechargeable sur d’autres accélérateurs. Les bibliothèques Diffusers et Kernels sont aussi concernées. Plusieurs pipelines d’images, dont FLUX.1-dev, FLUX.2 et Qwen-Image, sont compatibles avec TorchNeuron en mode immédiat ou compilé, avec un soutien croissant du parallélisme tensoriel. Le volet Kernels permet de partager des opérations optimisées écrites avec NKI, avec des gains pouvant atteindre 2,4 fois pour certains transformers vision. L’initiative est présentée comme une avancée importante pour rapprocher le matériel AWS de l’écosystème ouvert Hugging Face, tout en simplifiant le partage de modèles entre accélérateurs. Les équipes d’AWS sont remerciées pour leur travail sur TorchNeuron et sur l’intégration avec les outils Hugging Face.