Hugging Face launches a Space to check whether models can run on AWS Trainium
Hugging Face releases the torch-neuronx model support Space, which lets users view, by download count or trend, whether models on the Hub can run on AWS Trainium, and reports the number of sharded cores, compile-inclusive latency, and peak memory. Currently it only tests Transformers text-generation and image-text models, and only models that fit on 4 Trainium2 devices (up to 16 cores); each model generates 32 tokens for validation, comparing against CPU results in single-core mode, and in sharded mode a small model judges whether the English is fluent. Results come from the public dataset aws-neuron/torch-neuronx-model-support, and the Space also provides run scripts and a NEFF cache command that skips compilation.
Why read this
AI score 52
The original text provides the query entry point, testing criteria, and reproduction scripts, so readers can use them to determine whether a given model can run on Trainium.
Topics
Body
Look up a model
The Space opens on the Selected view: our own selection of models picked among the top models on the Hub by downloads and by trending score. Each row gives the result, the number of Trainium cores the model is sharded over, the wall time of the run (compilation included) and the peak device memory.
To look up a given model, switch to the Trending or Downloads view and type part of its name. These views follow the Hub's own rankings, one task at a time, whether a model was tested or not. When a model was left out on purpose, the Space tells you why (most of the time it was not prioritized, sometimes it is truly incompatible). You can also pick the versions of the Neuron compiler ( neuronx-cc ) and of Transformers to look at.
At the time of writing, only Transformers models are tested, and our selection is limited to text-generation and image-text-to-text models. We plan to expand to diffusers, sentence transformers, and other transformers tasks.
Run it yourself
Select a model to see how to run it. The Space lists the Neuron stack the model was tested on, the packages to install on top of it, and the run.py script that ran. The script follows the recipe of the TorchNeuron announcement : it loads the model with from_pretrained , sharded when it needs to be, then generates with a static cache and torch.compile(backend="neuron") .
On a Trainium2 ( trn2 ) instance with that stack, for example in the AWS Neuron Deep Learning Container , install the packages, then run:
No need for torchrun : the script launches its own torch.distributed ranks when the model is sharded.
A first run starts by compiling the model's graphs. To skip that step, copy the graphs the same compiler built during the test, from the public NEFF cache bucket. The Space gives you the command, with the right compiler version:
What "runs" means
Each test loads the model on a trn2 instance, sharded over the smallest number of Trainium cores that can hold its weights. Most models are sharded with tensor parallelism. Mixture-of-experts models without a tensor-parallel plan use expert parallelism only. The test then generates 32 tokens and checks them. On a single core, the tokens are compared with those of a CPU run of the same checkpoint. When the model is sharded, a small judge model checks that the text is fluent English.
At the time of writing, only models that fit on 4 Trainium2 devices are tested, which means 16 cores at most. A PASS says nothing about throughput, or about longer generations.
The results are a dataset
Everything the Space shows comes from the public aws-neuron/torch-neuronx-model-support dataset, with one Parquet file per pair of neuronx-cc and library versions.
Each row holds a model's result, the error when it failed, the versions of the whole stack and the run.py that ran. You can query it directly:
The dataset card gives the exact meaning of each result.
Next time you wonder whether a model runs on Trainium, look it up .
Datasets mentioned in this article 1
aws-neuron/torch-neuronx-model-support
Spaces mentioned in this article 1
torch-neuronx model support
Which Hub models run on AWS Trainium with torch-neuronx
More Articles from our Blog
TorchNeuron: running HF models on AWS Trainium is easier than ever
+3
Hugging Face ❤️ TorchTPU: TPUs, Natively
Datasets mentioned in this article 1
aws-neuron/torch-neuronx-model-support
Spaces mentioned in this article 1
torch-neuronx model support
Which Hub models run on AWS Trainium with torch-neuronx