Hugging Face ajoute une prise en charge native de PyTorch pour les puces AWS Trainium

Hugging Face a présenté TorchNeuron, un nouveau backend natif qui permet aux modèles PyTorch de s'exécuter directement sur le matériel AWS Trainium. Lancée en version bêta publique le 20 octobre 2026, cette mise à jour supprime la nécessité d'utiliser des flux de travail de compilation spécialisés, permettant aux développeurs d'utiliser les API PyTorch standard pour l'entraînement et l'inférence sur les accélérateurs d'IA d'Amazon.

Ce qui s'est passé

Pendant plusieurs années, l'exécution des modèles Hugging Face sur les puces AWS Trainium et Inferentia nécessitait l'utilisation d'Optimum Neuron. Cette bibliothèque servait de pont entre PyTorch et le matériel, reposant sur le paquet torch-neuronx construit sur PyTorch XLA. Bien que fonctionnelle, cette approche obligeait les développeurs à adopter une mentalité centrée sur le compilateur. Ils devaient gérer des tenseurs paresseux (lazy tensors) au lieu d'une exécution immédiate, se limiter à des formes statiques et gérer des recompilations fréquentes chaque fois que le graphe de calcul changeait. Optimum Neuron maintenait du code personnalisé, y compris des classes de modélisation et des entraîneurs spécifiques, pour masquer cette complexité aux utilisateurs.

TorchNeuron change cette architecture en s'intégrant directement dans PyTorch lui-même. Il expose neuron comme un périphérique natif, similaire à la façon dont les développeurs interagissent avec les GPU CUDA ou les périphériques CPU. Ce changement permet l'exécution immédiate (eager execution), ce qui signifie que les opérations sont effectuées dès qu'elles sont appelées, plutôt que d'être mises en file d'attente pour une compilation ultérieure. Il prend également en charge les API distribuées natives telles que Fully Sharded Data Parallel (FSDP) et Tensor Parallelism (TP), ainsi que les primitives distribuées propres à PyTorch comme DTensor et FSDP2. De plus, il fonctionne parfaitement avec torch.compile, permettant des optimisations de performance sans quitter le flux de travail PyTorch standard.

Par conséquent, les principales bibliothèques Hugging Face, notamment Transformers, Diffusers, Accelerate, PEFT et TRL, s'exécutent désormais directement sur Trainium. Les développeurs peuvent utiliser les mêmes classes et scripts qu'ils utiliseraient sur n'importe quel autre accélérateur. Hugging Face recommande désormais aux utilisateurs de migrer depuis Optimum Neuron et d'utiliser ces bibliothèques directement avec le nouveau backend. La mise à jour vise à simplifier l'expérience développeur et à faciliter le partage de modèles sur le Hugging Face Hub pour ceux qui utilisent l'infrastructure AWS.

Comment cela fonctionne

Le mécanisme central derrière TorchNeuron est son intégration en tant que backend PyTorch natif. Au lieu de traduire les opérations PyTorch via une couche de compilateur externe comme XLA, TorchNeuron permet à PyTorch de reconnaître neuron comme un périphérique de première classe. Cela active l'exécution en mode eager, où chaque opération est exécutée immédiatement sur le matériel. Cela contraste avec le modèle précédent d'évaluation paresseuse, qui nécessitait de construire un graphe de calcul complet avant que l'exécution ne puisse commencer.

En prenant en charge les primitives distribuées natives, TorchNeuron exploite les outils intégrés de PyTorch pour mettre à l'échelle les modèles sur plusieurs puces. Des fonctionnalités comme FSDP2 permettent une utilisation efficace de la mémoire pendant l'entraînement en fragmentant les paramètres du modèle sur les périphériques. Pour l'inférence, le backend prend en charge le parallélisme tensoriel, qui divise les couches du modèle sur plusieurs accélérateurs pour accélérer le traitement. Le système encourage également l'utilisation de caches statiques pour les tâches de génération afin de minimiser la surcharge de recompilation, garantissant que les performances restent élevées même avec la flexibilité de l'exécution immédiate.

Détails clés

• TorchNeuron est disponible en version bêta publique à partir du 20 octobre 2026.
• Les développeurs peuvent installer le paquet en utilisant la commande pip install torch-neuronx.
• Le nouveau backend prend en charge l'exécution immédiate, torch.compile et les API distribuées natives comme FSDP et TP.
• Dans Transformers, 89 % des 100 architectures de modèles de génération de texte les plus utilisées sont prises en charge nativement, contre 46 % avec Optimum Neuron v0.4.5.
• Les utilisateurs peuvent déplacer les modèles vers les périphériques Neuron en passant l'argument device_map=neuron dans leur code.
• Hugging Face recommande d'utiliser directement Transformers, Diffusers et d'autres bibliothèques plutôt que de passer par Optimum Neuron.

Pourquoi c'est important

Cette mise à jour abaisse considérablement la barrière à l'entrée pour les ingénieurs souhaitant utiliser AWS Trainium. Auparavant, la nécessité de penser comme un compilateur et de gérer des formes statiques créait une courbe d'apprentissage abrupte. En alignant l'expérience Trainium sur les flux de travail PyTorch standard, Hugging Face permet aux équipes de réutiliser leur code et leur expertise existants. Cela réduit l'effort d'ingénierie nécessaire pour porter les modèles sur le matériel AWS et minimise le risque de bugs associés aux étapes de compilation personnalisées.

L'augmentation du nombre d'architectures de modèles prises en charge élargit également l'utilité de Trainium pour les charges de travail de production. Avec 89 % des meilleurs modèles de génération de texte pris en charge nativement, les organisations peuvent déployer une gamme plus large de modèles de pointe sans attendre des intégrations personnalisées. Cette parité avec d'autres accélérateurs signifie que le choix du matériel peut être davantage dicté par le coût et la disponibilité plutôt que par des contraintes de compatibilité logicielle.

Ce que vous pouvez faire

• Installez le nouveau backend en exécutant pip install torch-neuronx dans votre environnement.
• Mettez à jour vos scripts d'inférence pour utiliser device_map=neuron lors du chargement des modèles avec Transformers.
• Remplacez les dépendances Optimum Neuron par des appels directs aux bibliothèques Hugging Face comme Transformers et Diffusers.
• Implémentez des caches statiques pour les tâches de génération de texte afin d'éviter les recompilations fréquentes et d'améliorer les performances.
• Testez vos configurations d'entraînement distribué existantes en utilisant les primitives PyTorch natives comme FSDP2 et DTensor sur des instances Trainium.
• Consultez le blog PyTorch pour obtenir des détails techniques supplémentaires sur la manière dont TorchNeuron s'intègre à l'écosystème PyTorch plus large.