Votre modèle tourne-t-il sur AWS Trainium ? Vérifiez-le

Dans un précédent article, nous avons présenté TorchNeuron, le backend PyTorch natif pour AWS Trainium (le torch-neuronx package). Avec lui, Transformers exécute son propre code de modélisation sur Trainium, comme sur tout autre appareil, donc tout modèle Transformers du Hub peut s'exécuter. Lesquels le font réellement ?

L'Space torch-neuronx model support répond à cette question pour les modèles les plus populaires du Hub. Il vous indique si un modèle s'exécute sans modification sur Trainium, et vous donne le script pour l'exécuter vous-même.

Rechercher un modèle

L'Space s'ouvre sur la vue Selected : notre propre sélection de modèles choisis parmi les meilleurs modèles du Hub par téléchargements et par score de tendance. Chaque ligne indique le résultat, le nombre de cœurs Trainium sur lesquels le modèle est shardé, le temps réel d'exécution (compilation incluse) et la mémoire maximale de l'appareil.

Pour rechercher un modèle donné, passez à la vue Trending ou Downloads et tapez une partie de son nom. Ces vues suivent les classements du Hub lui-même, une tâche à la fois, qu'un modèle ait été testé ou non. Lorsqu'un modèle a été écarté délibérément, l'Space vous dit pourquoi (la plupart du temps, il n'a pas été priorisé, parfois il est réellement incompatible). Vous pouvez également choisir les versions du compilateur Neuron (neuronx-cc) et de Transformers à consulter.

Au moment de la rédaction, seuls les modèles Transformers sont testés, et notre sélection est limitée aux modèles text-generation et image-text-to-text . Nous prévoyons de nous étendre aux diffusers, sentence transformers et autres tâches de transformers.

Exécutez-le vous-même

Sélectionnez un modèle pour voir comment l'exécuter. L'Space liste la pile Neuron sur laquelle le modèle a été testé, les packages à installer par-dessus, et le run.py script qui a été exécuté. Le script suit la recette de l' annonce TorchNeuron: il charge le modèle avec from_pretrained, shardé lorsque nécessaire, puis génère avec un cache statique et torch.compile(backend="neuron").

Sur une instance Trainium2 (trn2) avec cette pile, par exemple dans le conteneur AWS Neuron Deep Learning Container, installez les packages, puis exécutez :

python3 run.py

Pas besoin de torchrun: le script lance son propre torch.distributed rang lorsque le modèle est shardingé.

Une première exécution commence par compiler les graphes du modèle. Pour sauter cette étape, copiez les graphes que le même compilateur a construits pendant le test, depuis le bucket public cache NEFF . Le Space vous donne la commande, avec la bonne version du compilateur :

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

Ce que « runs » signifie

Chaque test charge le modèle sur une instance trn2 , shardingé sur le plus petit nombre de cœurs Trainium pouvant contenir ses poids. La plupart des modèles sont shardingés avec le parallélisme de tenseurs. Les modèles mixture-of-experts sans plan de parallélisme de tenseurs utilisent uniquement le parallélisme d'experts. Le test génère ensuite 32 tokens et les vérifie. Sur un seul cœur, les tokens sont comparés à ceux d'une exécution CPU du même checkpoint. Lorsque le modèle est shardingé, un petit modèle juge vérifie que le texte est un anglais fluide.

Au moment de la rédaction, seuls les modèles qui tiennent sur 4 appareils Trainium2 sont testés, soit 16 cœurs au maximum. Un PASS ne dit rien du débit, ni des générations plus longues.

Les résultats sont un jeu de données

Tout ce que le Space affiche provient du jeu de données public aws-neuron/torch-neuronx-model-support , avec un fichier Parquet par paire de neuronx-cc et de versions de bibliothèque.

Chaque ligne contient le résultat d'un modèle, l'erreur en cas d'échec, les versions de toute la pile et le run.py qui a tourné. Vous pouvez l'interroger directement :

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"]]

La fiche du jeu de données donne la signification exacte de chaque résultat.

La prochaine fois que vous vous demanderez si un modèle tourne sur Trainium, allez le vérifier.