Tools

AI-generated text

Hugging Face provides searchable results for which Transformers models run on AWS Trainium

Hugging Face published a Space that reports which Transformers models on the Hub run unmodified on AWS Trainium using the torch-neuronx backend.

Hugging Face provides searchable results for which Transformers models run on AWS Trainium

Hugging Face has published a Space that reports which popular Transformers models on the Hub run unmodified on AWS Trainium using the torch-neuronx (TorchNeuron) native PyTorch backend. The interface indicates whether a model ran on Trainium, how many Trainium cores it was sharded across, the wall time of the run (including compilation) and the peak device memory.

What you can search and how the Space works

The Space opens on a “Selected” view that lists Hugging Face’s own selection of top models by downloads and trending score. Users can switch to “Trending” or “Downloads” views and type part of a model name; those views follow the Hub’s rankings per task. If a model was deliberately omitted from testing, the Space explains why (usually low priority, sometimes actual incompatibility).

You can also pick the Neuron compiler (neuronx-cc) and Transformers library versions for the tests.

Which models and tasks are tested today

At the time of writing, only Transformers models are tested, and the selection is limited to text-generation and image-text-to-text tasks. Hugging Face plans to expand testing to diffusers, sentence transformers and other Transformers tasks.

How to run a tested model yourself

Selecting a model shows the Neuron stack used for the test, the packages to install, and the run.py script that was executed. The script follows the recipe from the TorchNeuron announcement: it loads the model with from_pretrained, shards it when necessary, then generates using a static cache and torch.compile(backend="neuron").

On a Trainium2 (trn2) instance with that stack (for example the AWS Neuron Deep Learning Container), install the required packages and run:

python3 run.py

There is no need to use torchrun: the script launches its own torch.distributed ranks when the model is sharded. The first run compiles the model graphs; to skip compilation you can copy the graphs that the same compiler produced during the test from the public NEFF cache bucket. The Space provides the exact command and the matching compiler version, for example:

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

What “runs” means in the tests

Each test runs on a trn2 instance and shards the model across the smallest number of Trainium cores that can hold its weights. Most models use tensor parallelism; mixture-of-experts models without a tensor-parallel plan use expert parallelism only. The test generates 32 tokens and checks the outputs: on a single core the tokens are compared to a CPU run of the same checkpoint, while for sharded runs a small judge model checks the text is fluent English.

As of this writing, only models that fit on 4 Trainium2 devices are tested (up to 16 cores). A PASS result does not indicate throughput or behavior for longer generations.

The results are published as a dataset

All data shown in the Space come from the public aws-neuron/torch-neuronx-model-support dataset, with one Parquet file per neuronx-cc and library version pair. Each row records a model’s result, the error when it failed, the versions of the full stack and the run.py that was executed. The dataset can be queried directly, for example with pandas:

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

The dataset card explains the exact meaning of each result.

Conclusion

If you want to know whether a particular Hub Transformers model runs on AWS Trainium with torch-neuronx, consult the Hugging Face Space: it provides per-model test results, the run script and commands, and the underlying public dataset for inspection.