Does Your Model Run on AWS Trainium? Look It Up

Does Your Model Run on AWS Trainium? Look It Up

In a previous post, TorchNeuron, the native PyTorch backend for AWS Trainium provided by the torch-neuronx package, was introduced. With this backend, Transformers uses its regular modeling code on Trainium, as it does on other devices. This means that any Transformers model on the Hub may be able to run on Trainium, but actual compatibility varies.

The torch-neuronx model support Space reports compatibility for many of the most popular models on the Hub. It indicates whether a model runs without modification on Trainium and provides the script used to run it.

Look up a model

The Space opens on the Selected view, which contains models selected from the Hub's top models based on download counts and trending scores. Each row shows the test result, the number of Trainium cores across which the model is sharded, the wall-clock runtime including compilation, and peak device memory usage.

To search for a particular model, switch to the Trending or Downloads view and enter part of its name. These views follow the Hub's rankings one task at a time, regardless of whether a model has been tested. When a model was intentionally omitted, the Space explains the reason. Usually, the model was not prioritized, although some models are incompatible.

The Space also lets you select the versions of the Neuron compiler, neuronx-cc, and Transformers to inspect.

At the time of writing, testing covers only Transformers models from the text-generation and image-text-to-text tasks. Planned expansions include diffusers, sentence transformers, and additional Transformers tasks.

Run it yourself

Select a model to view the instructions for running it. The Space lists the Neuron software stack used for testing, the packages that must be installed, 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, and generates text using a static cache and torch.compile(backend="neuron").

On a Trainium2 (trn2) instance with the same stack, such as an AWS Neuron Deep Learning Container, install the required packages and run:

python3 run.py

torchrun is not required. When the model is sharded, the script launches its own torch.distributed ranks.

The first run compiles the model's graphs. That step can be skipped by copying the graphs produced by the same compiler during testing from the public NEFF cache bucket. The Space provides the command with the appropriate compiler version:

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

What “runs” means

Each test loads a model on a trn2 instance and shards it 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 then generates 32 tokens and checks the output. For a model running on a single core, the generated tokens are compared with those from a CPU run using the same checkpoint. For sharded models, a small judge model checks whether the generated text is fluent English.

At the time of writing, testing is limited to models that fit on four Trainium2 devices, or 16 cores at most. A PASS result does not provide information about throughput or longer generations.

The results are a dataset

All results displayed by the Space come from the public aws-neuron/torch-neuronx-model-support dataset. It contains one Parquet file for each pair of neuronx-cc and library versions.

Each row includes the model's result, the error reported when a test failed, the versions of the complete software stack, and the run.py script that was executed. The data can be queried directly:

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 the meaning of each result.