Contexte et objectif du Space

Hugging Face a publié le Space torch-neuronx-model-support pour répondre à la question : quel modèle du Hub fonctionne sans modification sur les accélérateurs AWS Trainium ? Le tableau répertorie les modèles les plus téléchargés ou les plus tendance, indique le nombre de cœurs Trainium utilisés, le temps d’exécution incluant la compilation et la consommation maximale de mémoire. Cette visibilité permet aux ingénieurs de choisir rapidement un modèle compatible avant de lancer un déploiement.

Architecture du backend TorchNeuronX

Le package torch-neuronx agit comme un backend natif de PyTorch. Lors de l’exécution, le script run.py charge le modèle via from_pretrained, applique le sharding nécessaire, puis lance la génération avec torch.compile(backend="neuron"). Sur une instance Trainium2 (trn2) le modèle est réparti sur le plus petit nombre de cœurs capable de contenir ses poids ; la plupart des modèles utilisent le parallélisme tensoriel, tandis que les Mixture‑of‑Experts s’appuient uniquement sur le parallélisme d’experts. Le processus de compilation crée des graphes NEFF stockés dans un bucket public, réutilisables via la commande suivante :

hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc- /tmp/neff_cache

Cette étape évite la recompilation lors de nouvelles exécutions, réduisant le temps de démarrage.

Méthodologie de test et limites observées

Chaque test charge le modèle sur une instance trn2, le shard sur le nombre minimal de cœurs (maximum 16 cœurs, soit 4 appareils Trainium2). Le benchmark génère 32 tokens et compare le résultat avec une exécution CPU du même checkpoint ; en mode sharding, un modèle de jugement vérifie la fluidité de l’anglais produit. Les résultats sont stockés dans le dataset public aws-neuron/torch-neuronx-model-support, un fichier Parquet par combinaison de version du compilateur neuronx-cc et de la bibliothèque Transformers. La requête suivante illustre l’accès aux données :

import pandas as pd
df = pd.read_parquet("hf://datasets/aws-neuron/torch-neuronx-model-support/data")
df[df.model_id == "Qwen/Qwen3-0.6B"][["neuronx_cc", "library", "library_version", "neuron_cores", "result"]]

Les tests se limitent aux modèles qui tiennent sur 4 appareils (16 cœurs). Un statut PASS ne garantit ni le débit ni la stabilité sur des générations plus longues, et les modèles non testés sont souvent exclus par manque de priorité ou d’incompatibilité réelle.

Implications pour les développeurs

Le Space fournit une source unique pour valider la compatibilité avant d’investir dans une instance Trainium. En copiant les graphes NEFF pré‑compilés, les équipes réduisent le coût de compilation et peuvent lancer python3 run.py directement, sans recourir à torchrun. La visibilité sur le nombre de cœurs requis aide à dimensionner les ressources cloud et à anticiper les coûts d’infrastructure. La feuille de route indique l’ajout futur de diffusers, de sentence‑transformers et d’autres tâches, élargissant ainsi le champ d’application du backend TorchNeuronX.