Les notes de version de Neuron 2.33 d’AWS situent la bêta publique de TorchNeuron au 19 octobre. L’annonce de PyTorch, également datée du 19 octobre, indique que le lancement de la bêta publique et de la version open source a lieu le 20 octobre. Le paquet torch-neuronx enregistre dans PyTorch un périphérique neuron pour Trainium2 et Trainium3, permettant aux opérations de modèle et de tenseur prises en charge d’utiliser les API PyTorch habituelles. Pour les ingénieurs qui envisagent ces puces, le changement utile est un chemin plus court entre un programme PyTorch existant et un essai sur le périphérique. Le bon fonctionnement et l’efficacité de ce programme dépendent toujours de ses opérations, de ses formes de tenseurs et de son environnement logiciel.

L’ancienne voie PyTorch pour Trainium utilisait la bibliothèque torch-neuronx fondée sur XLA et du code de périphérique torch_xla. Selon AWS, le backend natif utilise l’extension PrivateUse1 de PyTorch et remplace xm.xla_device() par un placement sur le périphérique neuron. C’est une modification importante de l’interface pour les projets qui utilisent déjà Neuron, même si le nom du paquet reste familier. Cette version s’applique à l’entraînement et à l’inférence sur les instances Trn2 et Trn3 ; la matrice de compatibilité ne mentionne pas Trn1.

Ce qui change dans un programme PyTorch

Les notes de version d’AWS montrent un modèle déplacé avec .to("neuron") et un tenseur d’entrée créé sur device="neuron". Les opérations prises en charge de torch.nn, d’autograd et des optimiseurs s’exécutent ensuite via le backend. Le code écrit avec l’API torch.accelerator de PyTorch, indépendante du périphérique, peut détecter Neuron sans réécriture propre au périphérique, selon la présentation de TorchNeuron.

Deux modes d’exécution sont disponibles. Le mode eager exécute les opérations PyTorch à mesure que Python les appelle, en conservant le flux de contrôle ordinaire, les affichages et les points d’arrêt. Chaque opération est compilée lors de sa première utilisation, puis les appels correspondants peuvent utiliser un cache. Les ingénieurs peuvent aussi compiler un modèle ou une région sélectionnée avec torch.compile(model, backend="neuron", dynamic=False). Le backend natif compile les graphes de calcul avant et arrière pris en charge ; le code extérieur à la région capturée continue de s’exécuter en mode eager. L’annonce de PyTorch explique aussi comment définir torch.compiler.set_default_backend("neuron") une fois avant la compilation, pour que les appels de compilation ultérieurs puissent omettre l’argument du backend.

Cette version comprend un backend neuron pour torch.distributed, avec des opérations collectives et point à point documentées. AWS cite DistributedDataParallel, FSDP2, DTensor et DeviceMesh, ainsi que le parallélisme de pipeline, parmi les API prises en charge. Ce sont des pistes utiles à examiner pour un projet d’entraînement distribué, mais les notes de version ne démontrent pas qu’un traitement multinœud existant s’exécutera sans modification ni qu’il atteindra un objectif précis de débit.

L’évaluation commence par un environnement compatible

Le guide d’installation rapide d’AWS exige une instance Trn2 ou Trn3 et soit une Neuron Base DLAMI, soit une AMI Ubuntu prise en charge. La matrice de cette version indique Linux x86_64 sur Ubuntu 24.04 ou 26.04. Ubuntu 22.04 n’est plus pris en charge pour cette version. Elle mentionne Python 3.12, 3.13 et 3.14, avec torch-neuronx 2.13.3.0 pour PyTorch 2.13.0 ou torch-neuronx 2.14.3.0 pour PyTorch 2.14.0. La famille du paquet doit correspondre à la version de PyTorch.

Le guide rapide installe torch-neuronx==2.14.* ainsi que neuronx-cc==2.* et nki depuis les index de paquets Neuron et PyTorch. Il vérifie ensuite la disponibilité du périphérique et compare une petite opération matricielle sur Neuron avec un résultat obtenu sur CPU. Cela vérifie le chemin vers le périphérique installé dans l’exemple d’AWS. Cela ne valide pas le modèle du lecteur. Une évaluation suivante raisonnable consiste à exécuter les calculs avant et arrière du modèle réel, à examiner les replis sur CPU enregistrés avec torch_neuronx.get_fallback_ops(), à comparer les sorties avec celles de l’environnement actuel et à mesurer séparément la compilation et les exécutions en régime établi. Ces dernières étapes sont les questions d’évaluation proposées par BIG CHANGE, et non des tests que nous avons effectués.

Ce que la bêta peut encore changer dans le travail

En mode eager, une opération non prise en charge peut se replier sur le CPU et émettre un avertissement. Avec torch.compile, une opération capturée dans un graphe que le backend Neuron ne peut pas convertir fait échouer la compilation par défaut. AWS renvoie les développeurs à sa table de compatibilité des opérateurs pour connaître la couverture de chaque opération. Une exécution eager qui aboutit ne prouve donc pas à elle seule que tout s’exécute sur le périphérique ni que le modèle est compatible avec la compilation.

Les notes de version précisent aussi que torch.compile exige dynamic=False ; une nouvelle forme d’entrée peut imposer une nouvelle compilation. AWS suggère le padding ou le regroupement par formes, appelé bucketing, pour réduire leur variation. La première exécution inclut la compilation lorsqu’aucun artefact correspondant n’est en cache. Ces détails comptent pour chronométrer une charge de travail : la latence de compilation et la répartition des formes d’entrée peuvent modifier le résultat d’une évaluation.

D’autres limites explicites restreignent la compatibilité annoncée. La bêta ne prend pas en charge torch.sparse, torch.func, torch.nested ni torch.distributions. Seules certaines transformations torch.fft s’exécutent sur Neuron, et float64 s’exécute en float32 avec un avertissement. L’intégration de TorchAO et la prise en charge du type de données composite MXFP8 sont absentes de cette version. Une équipe qui utilise ces fonctions doit vérifier son propre chemin d’exécution avant de considérer le placement sur le périphérique comme un plan de migration.

Le grand changement

TorchNeuron place Trainium derrière les interfaces natives de périphérique et de compilation de PyTorch dans une bêta publique. Cela rend un premier essai plus direct pour les programmes PyTorch pris en charge et offre aux utilisateurs actuels de Neuron une autre voie que leur code fondé sur XLA. Les éléments décisifs pour l’adoption viendront de chaque charge de travail : ses opérateurs restent-ils sur le périphérique prévu, ses résultats numériques demeurent-ils acceptables, et les performances et coûts mesurés justifient-ils l’instance et le travail de migration ? BIG CHANGE a examiné la documentation et les annonces pour cette analyse ; nous n’avons pas exécuté TorchNeuron ni vérifié indépendamment les affirmations d’AWS sur les performances.

Sources et lectures complémentaires