AWS ouvre TorchNeuron, un backend PyTorch natif pour Trainium, en bêta publique

AWS ouvre TorchNeuron, un backend PyTorch natif pour Trainium, en bêta publique

AWS a annoncé le 20 octobre que TorchNeuron, un backend PyTorch open source pour ses puces d'IA Trainium, était entré en bêta publique. Un article plus long d'ingénieurs d'AWS sur le blog de PyTorch est daté du 19 octobre et mis à jour le 20 octobre, jour de la PyTorch Conference North America. L'annonce d'AWS précise que le code est open source pour les instances Amazon EC2 Trainium2 et Trainium3. L'article indique la licence Apache 2.0 et un dépôt GitHub sous aws-neuron. Ce numéro du 27 octobre en rend compte une semaine après la date d'origine. Les deux pages présentent le point de vue d'AWS, mais l'article PyTorch cite le directeur technique de Deepgram : dans les tests évalués par son entreprise, Trainium traitait environ 60 % de secondes d'audio entrant en plus par dollar que ses plateformes GPU existantes. Il s'agit du témoignage d'un client, pas d'une mesure indépendante ni d'une garantie générale des performances de TorchNeuron.

Ce que cela change pour les développeurs

Jusqu'ici, exécuter du code PyTorch sur Trainium passait par l'outillage distinct d'AWS. TorchNeuron enregistre Trainium comme un appareil PyTorch ordinaire grâce à PrivateUse1, le point d'extension officiel de PyTorch pour le matériel personnalisé, et à l'interface `torch.accelerator`. Dans le cas le plus simple, un développeur déplace un modèle avec `.to('neuron')` et conserve le reste de sa boucle d'entraînement. Deux modes d'exécution sont pris en charge. Le mode eager exécute les opérations une à une, ce qui convient au débogage, et l'article ajoute un mode expérimental « adaptive eager » destiné à réduire le coût de lancement et à fusionner des opérations sans modifier le code. Le mode `torch.compile` capture le modèle sous forme de graphe et le confie au compilateur Neuron ; AWS le destine à la production. L'article nomme le chemin de compilation : TorchDynamo, AOTAutograd, Torch-MLIR, StableHLO, puis le compilateur Neuron, qui produit un fichier exécuté sur la puce.

Périmètre et limites

AWS cite une large couverture d'opérateurs pour le texte, la vision, la vidéo et d'autres modèles, la prise en charge des types BF16, FP32 et FP16, et les API distribuées standard : DDP, FSDP2, parallélisme tensoriel, parallélisme en pipeline, DTensor et DeviceMesh. La communication passe par les cœurs de communication collective de Trainium à l'intérieur d'un nœud et par l'Elastic Fabric Adapter entre nœuds. L'article est franc sur les lacunes. TorchAO et FP8 sont en bêta privée. Les flux non par défaut partagent une seule file matérielle : ils conservent l'ordre, mais ne disposent pas chacun d'un moteur dédié. Les erreurs apparaissent au point de synchronisation suivant plutôt qu'à l'opération fautive, et TorchNeuron relève la plus ancienne avec son contexte. Les noyaux personnalisés s'écrivent avec la Neuron Kernel Interface, et AWS fournit une bibliothèque open source de tels noyaux, dont des routines d'attention et de mélange d'experts. Une intégration du profileur relie les événements du framework à l'activité matérielle.

Comment lire l'annonce

La présentation du produit met la compatibilité avec PyTorch au premier plan. AWS dit que les développeurs conservent les API de PyTorch tout en accédant à Trainium, et un directeur d'ingénierie de Meta est cité pour saluer cette prise en charge. L'objectif du mode adaptive eager, se rapprocher des performances des graphes compilés, est un but énoncé dans l'article. L'article PyTorch donne bien le chiffre de Deepgram sur le coût par seconde d'audio entrant, sans vérification indépendante ni conclusion générale sur la vitesse. Il décrit aussi un balayage de 100 000 modèles Hugging Face dans des configurations réduites mais représentatives pour vérifier le chargement, l'exécution, la compilation et le placement des opérateurs : il ne s'agit pas de 100 000 comparaisons de performance avec des GPU. Une équipe qui envisage Trainium devrait faire passer son propre modèle par les modes eager et compilé, vérifier que les opérateurs dont elle a besoin sont couverts, et lire les notes de la bêta avant de migrer du travail de production.