Find Out Whether Your Model Runs on AWS Trainium
A torch-neuronx Space lists Transformers models tested on AWS Trainium, showing PASS results, cores, time, memory, scripts, and steps to rerun tests.
Key Takeaways
- The torch-neuronx model support Space shows which popular Transformers models run unmodified on AWS Trainium, including test results, core counts, run time, memory use, and reproduction scripts.
- A “PASS” means the model ran on a Trainium2 instance and produced checked output; it is a compatibility check, not a throughput or production benchmark.
- The Space covers text-generation and image-text-to-text Transformers models at present, with plans to expand to other model types and tasks.
- Users can reproduce a tested model on Trainium2 by using the listed Neuron stack, installing required packages, and running the provided run.py script; compiled graphs can be reused from the NEFF cache to avoid waiting.
- The test results are published as an open dataset, making it easier for developers to query compatibility status, versions, errors, and scripts.
Meta description: The torch-neuronx model support Space shows which Transformers models run on Trainium, with core counts, run times, memory and test scripts.
Tags: Trainium, torch-neuronx, Transformers, AWS Neuron, model support
Featured image: abstract illustration of model cards and Trainium cores, friendly style, no people, no logos, no text
At the time of writing, the torch-neuronx model support Space tells developers whether popular Transformers models from the Hub run unmodified on AWS Trainium, and it gives the script used to reproduce each test. It follows the earlier introduction of TorchNeuron, the native PyTorch backend for AWS Trainium, packaged as torch-neuronx. With that backend, Transformers runs its own modeling code on Trainium, as it does on other devices, so many Hub models may be candidates. The Space turns that broad possibility into a searchable list of tested results.
What the Space shows first
The Space opens in the Selected view. That view contains a curated group of models chosen from the Hub's top models 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 including compilation, and the peak device memory. This structure helps users see both the outcome and the resources used to reach it.
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 why; most of the time it was not prioritized, and sometimes it is truly incompatible. Users can also pick the versions of the Neuron compiler, neuronx-cc, and of Transformers.
At the time of writing, only Transformers models are tested, and the selection is limited to text-generation and image-text-to-text models. The Space plans to expand to diffusers, sentence transformers, and other transformers tasks.
What a PASS measures
A PASS in the Space means a model ran on a Trainium2 instance, shown as trn2, and produced checked output. For each test, the model is loaded onto a trn2 instance and spread across the fewest Trainium cores needed to hold its weights. Most models are sharded using tensor parallelism. When a mixture-of-experts model has no tensor-parallel plan, the Space uses expert parallelism only.
The run then produces 32 tokens, and those tokens are checked. For a single-core run, the Space compares the tokens with the output from a CPU run of the same checkpoint. For sharded runs, a small judge model checks whether the output is fluent English.
The result is a functional check, not a performance benchmark. Current coverage stops at models that fit on 4 Trainium2 devices, or 16 cores at most. A PASS does not show throughput or predict behavior for longer generations.
How to rerun a test on Trainium2
Selecting a model shows 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 run.py uses the same setup described in the TorchNeuron announcement. It calls from_pretrained to load the model, applies sharding when required, and generates with a static cache under torch.compile(backend="neuron").
Users can run the same script on a Trainium2 instance, shown as trn2, that has the listed stack. An example environment is the AWS Neuron Deep Learning Container. After installing the listed packages, run the script. No need for torchrun: the script launches its own torch.distributed ranks when the model is sharded.
The first execution builds the model graphs through compilation. If users want to avoid that wait, they can reuse the graphs created by the same compiler during the Space test. Those graphs are available from the public NEFF cache bucket, and the Space provides the matching command for the compiler version.
The test results are open as a dataset
The interface is built on the public aws-neuron/torch-neuronx-model-support dataset. The dataset stores one Parquet file for each pair of neuronx-cc and library versions.
Each record contains the test outcome, the error if the run failed, the versions that made up the stack, and the run.py used. Developers can query the dataset directly. The dataset card gives the exact meaning of each result.
For readers, the main value is a faster way to decide whether a model is worth trying on Trainium. It does not replace throughput testing or production validation. The next expansion should add more model types, while the current 32-token tests remain a quick compatibility check for popular Transformers workloads on Trainium. The Space is meant to be looked up whenever Trainium compatibility is in question.
Comments (0)