TorchNeuron : Amazon ouvre Trainium à PyTorch avec un simple changement de configuration

La fin de la réécriture obligatoire

Amazon Web Services a présenté TorchNeuron, une bibliothèque open source sous licence Apache 2.0 qui permet d'exécuter du code PyTorch directement sur ses accélérateurs Trainium2 et Trainium3. L'objectif principal est d'éliminer la barrière technique qui a historiquement empêché les développeurs de migrer leurs modèles vers du matériel alternatif : la nécessité de réécrire des applications complètes pour s'adapter à une nouvelle architecture.

Pendant plus d'une décennie, l'écosystème de NVIDIA s'est consolidé grâce à CUDA, un logiciel permettant aux applications PyTorch de fonctionner de manière native sur ses GPU. TorchNeuron attaque directement cet avantage concurrentiel en s'intégrant aux points d'extension officiels de PyTorch, spécifiquement via le backend PrivateUse1. Cela signifie que les développeurs peuvent conserver leur flux de travail habituel et simplement changer le périphérique cible via une ligne de configuration.

Architecture technique et modes d'exécution

TorchNeuron n'est pas un simple wrapper, mais un backend natif qui s'intègre profondément au compilateur de PyTorch. Il prend en charge deux modes d'exécution principaux : le mode eager, idéal pour la recherche et le débogage, et le mode torch.compile, optimisé pour les performances en production. En mode eager, chaque opération est compilée individuellement dans un fichier NEFF (Neuron Executable File Format) et exécutée sur le périphérique.

Pour les charges de travail à forte intensité, l'utilisation de torch.compile permet de capturer le modèle sous forme de graphe que le compilateur Neuron optimise globalement. Ce processus utilise Torch-MLIR pour convertir le graphe en StableHLO avant de générer le code final pour le matériel. De plus, la bibliothèque prend en charge les API distribuées standard telles que FSDP, DDP et DTensor, facilitant l'entraînement de modèles à grande échelle sans modifier la logique de parallélisme.

Limites techniques de la version initiale

Bien que le lancement soit significatif, la version bêta publique de TorchNeuron présente des limites concrètes que les équipes doivent considérer avant de migrer des charges de travail critiques. Selon la documentation technique, la première version omet environ 22 000 tests de la suite PyTorch, principalement liés à des opérateurs qui n'ont pas encore d'implémentation native dans Neuron.

Parmi les fonctions non prises en charge dans ce premier lancement figurent les formes d'entrée dynamiques (dynamic shapes), les tenseurs creux et certaines opérations d'autograd en mode forward. Ces omissions signifient que tout le code PyTorch ne sera pas immédiatement compatible, et les développeurs devront valider leurs modèles spécifiques pour s'assurer qu'ils ne dépendent pas de ces fonctionnalités exclues.

Impact stratégique et concurrence avec NVIDIA

Le mouvement d'Amazon n'est pas seulement technique, mais stratégique. En ouvrant TorchNeuron et en contribuant au projet principal de PyTorch pour rendre le framework plus agnostique par rapport au matériel, AWS cherche à éroder l'avantage de NVIDIA dans le logiciel. L'idée est que les développeurs choisissent le matériel en fonction du prix, des performances et de la disponibilité, et non de la compatibilité forcée par le code.

Colin Brace, VP de la technologie logicielle chez AWS Annapurna Labs, a indiqué que l'objectif est d'éliminer le 'lock-in' qui empêchait les développeurs d'avoir un choix réel. Avec plus de 30 clients initiaux testant la technologie, y compris des entreprises d'IA vocale comme Deepgram et Cartesia, Amazon démontre qu'il est possible d'obtenir des améliorations de l'efficacité des coûts sans sacrifier la qualité du modèle.

Comment commencer avec TorchNeuron

Pour les développeurs intéressés par la mise en œuvre de TorchNeuron, le processus d'installation est simple. La bibliothèque est disponible sur GitHub sous le dépôt aws-neuron/torch-neuronx. Une fois installée, la bascule vers Trainium s'effectue en modifiant le périphérique cible du modèle.

Il est important de noter que TorchNeuron nécessite des instances EC2 Trn2 ou Trn3. Les équipes qui utilisent déjà PyTorch peuvent commencer avec de petits modèles pour valider la compatibilité avant de passer à des charges de travail d'entraînement ou d'inférence à grande échelle. L'intégration avec Hugging Face Transformers permet également de charger des modèles pré-entraînés directement sur le périphérique Neuron.

import torch

# Définir un modèle simple
model = torch.nn.Linear(10, 1)

# Migrer vers Trainium (changement d'une ligne)
model = model.to('neuron')

# Exécuter une inférence
input_data = torch.randn(1, 10, device='neuron')
output = model(input_data)

print(output)

Poursuivre l’exploration

Corrections et transparence

Si vous repérez une information incorrecte ou obsolète, envoyez-nous la page et la source à vérifier.