Does Your Model Run on AWS Trainium? Look It Up

What it claims

The torch-neuronx model support Space reports whether popular Hugging Face Transformers models run unmodified on AWS Trainium2 through TorchNeuron, the native PyTorch backend in the torch-neuronx package. Coverage is currently limited to text-generation and image-text-to-text models. The Space has three views. Selected is a curated pick from the Hub's top models by downloads and trending score. Trending and Downloads follow the Hub's per-task rankings and list models whether or not they were tested. Results can be filtered by neuronx-cc compiler version and Transformers version. When a model was skipped, the Space states why: usually it was not prioritized, sometimes it is incompatible.

How it works

Each test loads the model on a trn2 instance over the smallest number of 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 script loads with from_pretrained, generates with a static cache and torch.compile(backend="neuron"), then checks 32 generated tokens. How those tokens are checked depends on sharding:

flowchart TD
    A[Load on smallest core count] --> B{Sharded?}
    B -- No, 1 core --> C[Compare tokens to CPU run of same checkpoint]
    B -- Yes --> D[Small judge model checks fluent English]
    C --> E[PASS / FAIL]
    D --> E

On a single core, the tokens must match a CPU run. On sharded runs, a judge model checks only that the text is fluent English, which is a weaker check.

How to use it

On a trn2 instance with the listed Neuron stack (for example, the AWS Neuron Deep Learning Container), install the packages the Space lists, then run:

python3 run.py

No torchrun is needed, because the script starts its own torch.distributed ranks when sharding. To skip first-run graph compilation, pull the graphs that the same compiler built during the test:

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

The dataset has one Parquet file per neuronx-cc and library version pair. Each row includes the error text for failures and the full stack versions:

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 defines each result value.

Limits

Only models that fit on 4 Trainium2 devices (16 cores) are tested. A PASS says nothing about throughput or about generations longer than 32 tokens. Reported wall time includes compilation. Diffusers, sentence transformers and other Transformers tasks are planned but not yet covered.