AWS Trainium Model Support: Check What Runs
Check whether popular Transformers models run on AWS Trainium, review tested versions and limits, and use the provided scripts to reproduce results yourself.
Does Your AI Model Run on AWS Trainium? A Practical Check
Just because a model shows up in Transformers doesn't mean it'll run cleanly on every accelerator. With AWS Trainium, the questions get more concrete: does your specific checkpoint work with the current software stack, how many Neuron cores does it consume, and can your team reproduce the run without getting stuck in compiler or compatibility debugging?
The torch-neuronx model support Space is meant to answer that kind of "will it work here?" question for a set of popular models on the Hugging Face Hub. It tells you whether a model passed a defined test, lists the software versions used, and includes a script you can run. That's helpful - but it still has to be read correctly. A passing test isn't a benchmark, and the tool's coverage isn't a guarantee that every model or workload will behave the same way.
This guide shows how to use the support Space, what its results do (and don't) establish, how to reproduce a test on Trainium2, and how to use the public dataset when you want to inspect results programmatically. The goal isn't to treat a compatibility check like a deployment decision. It's to make that decision with better evidence - and fewer surprises.
What the Trainium model support tool tells you
The TorchNeuron announcement introduced TorchNeuron, the native PyTorch backend for AWS Trainium, shipped via the torch-neuronx package. The idea is straightforward: Transformers can use its own model implementation on Trainium, similar to how it runs on other devices. That makes it plausible that a model on the Hub might work on Trainium - but it doesn't prove that it will work for every compiler and library combination.
The support Space narrows that down. It lists results for selected popular models and lets you search through Hub rankings like Trending or Downloads. For a tested model, each row can include the result, the number of Trainium cores used, elapsed wall time (including compilation), and peak device memory. If a model wasn't tested, the Space may also mark it as intentionally excluded and explain why (when it can).
That distinction matters. "Not tested" doesn't mean "doesn't run." It might just mean nobody tried that model in the defined scope. And a model appearing in a ranking doesn't mean it passed. In other words, the lookup is evidence about the tested combinations - not a blanket compatibility promise.
The software versions are part of the result
A model name alone won't tell you much. The Neuron compiler (neuronx-cc) and the Transformers version matter because they're inputs to the compatibility outcome. The Space lets you select specific compiler and Transformers versions so you can see whether a test applies to the stack you plan to use.
This is the part that often gets lost in compatibility discussions. If your production image pins different versions than the test environment, you should treat the result as a reference point, not as proof that your exact setup will behave the same way. When you evaluate a deployment, record the full software stack alongside the model checkpoint. Otherwise, changes over time can make an "it worked" result hard to reproduce.
Coverage has a defined scope
The source describes testing focused on Transformers models for text generation and image-text-to-text. The tool's selection was limited to models that fit on four Trainium2 devices - up to 16 Trainium cores. The project also mentioned plans to expand coverage to areas like Diffusers and Sentence Transformers, plus other Transformers tasks.
So treat this as a snapshot of what the tool tested, not a statement about every model family or future software release. If your workload is outside those tasks, or depends on a customized model implementation, you'll still need to test your workload directly.
How to look up and run a model
When you open the Space, you'll see a Selected view: a curated set chosen from popular Hub models based on downloads and trending score. It's fine for a quick scan, but it's not ideal if you're making an engineering decision about a specific checkpoint. To search for a particular model, switch to Trending or Downloads and enter part of its name. Those views follow Hub rankings per task, and they show whether each model has a tested result or not.
Use the view to answer a few early questions:
- Is the model listed, and does it have a test result?
- Which Neuron compiler and Transformers versions does that result correspond to?
- How many Trainium cores were used, and what memory and wall-time figures were recorded?
If the model isn't tested, don't treat an empty result like a failure - read the status. It may simply not have been prioritized. If it is tested, open the model entry to inspect the environment and the run script. That moves you from "someone says it works" to a concrete recipe you can examine and try.
Reproduce the test before making a deployment decision
For a selected model, the Space shows the Neuron stack used in testing, the packages installed on top of that stack, and the run.py script that was executed. The script follows the TorchNeuron recipe: it loads the checkpoint with from_pretrained, shards if needed, then generates using a static cache and torch.compile(backend="neuron").
On a Trainium2 (trn2) instance with the matching stack - for example, an AWS Neuron Deep Learning Container - the basic flow is:
python3 run.pyYou don't need to start it with torchrun. When sharding is required, the script handles its own torch.distributed ranks. If you add a second launcher on top of that, you can end up with confusion about process ownership and configuration.
Also, the first run compiles the model graphs. Since compilation time is included in the wall time reported by the test, the first launch won't represent steady-state generation performance. If you want to avoid recompiling the same graphs, the Space includes a command for copying the compiled NEFF files from a public cache bucket. The command format is:
hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cacheUse the exact compiler version tied to the result. A cache created for one compiler version isn't a drop-in replacement for another. More importantly: copying a cache is a convenience for reproducing the compatible test environment, not a substitute for validating your own runtime setup.
A practical evaluation sequence
For a real project, keep the process small and explicit:
- Find the exact model ID and check whether it has a tested result.
- Match the tested compiler and library versions to your target environment.
- Read the provided script and package list, not just the headline status.
- Run the script on the target instance type and verify it completes in your environment.
- Test your actual input shapes, generation lengths, and service requirements separately.
That way, you separate "can it run under a defined setup?" from "is it fast enough and reliable enough for our workload?" The first answers compatibility. The second requires your measurements.
What a passing test does - and does not - mean
The test procedure is intentionally bounded. Each model is loaded on a Trainium2 instance and sharded across the smallest number of Trainium cores that can fit its weights. Most models use tensor parallelism. Mixture-of-experts models without a tensor-parallel plan use expert parallelism only.
After loading, the test generates 32 tokens and checks the output. For a model running on a single core, the generated tokens are compared with a CPU run of the same checkpoint. For sharded models, a small judge model checks whether the generated text is fluent English. That's a practical compatibility signal, but it isn't the same as a full numerical-equivalence test or a full evaluation of model quality.
A PASS means a narrow thing: under the tested stack and setup, the model loaded and produced an output that met the test's check. It doesn't establish throughput, latency under concurrent requests, cost per token, or whether the output is good on your business task. And generating 32 tokens doesn't tell you how the model behaves during much longer generations.
Compatibility is one gate, not the whole deployment review
A deployment decision usually involves multiple gates:
- Can it load and execute? The support result is relevant here.
- Does it fit the target configuration? Compare the tested core count and memory figures to your instance and plan.
- Does it meet service requirements? Measure latency, throughput, concurrency, and generation length with representative traffic.
- Does the output meet product requirements? Evaluate quality, safety, and task-specific failure modes using your own data.
- Do the economics work? Compare total operating cost versus alternatives under the same workload.
When teams collapse everything into "the model runs," they end up with demos that technically work but don't meet production needs. A compatibility result reduces uncertainty at the first gate. It doesn't remove the need to measure the rest.
There's also a capacity boundary: the stated scope tested models that fit on four Trainium2 devices (up to 16 cores). If a bigger model falls outside that range, the dataset doesn't disprove it - it simply wasn't covered. The next step for those models is a targeted test, not an assumption based on a similarly named checkpoint.
How to work with the results dataset
The Space is a UI over a public dataset: aws-neuron/torch-neuronx-model-support. The dataset includes one Parquet file for each pair of neuronx-cc and library versions. Each row records the model result, the error when a test failed, versions for the software stack, the number of Neuron cores, and the run.py script used.
That structure is useful when you need more than a manual lookup. For example, you might compare results across library versions, inspect failures, or build a quick internal report for a set of candidate models. The dataset card defines what each status means, so check it before treating status values as interchangeable with your own terminology.
You can load the data with pandas and filter for a specific model ID:
import pandas as pd
df = pd.read_parquet(
"hf://datasets/aws-neuron/torch-neuronx-model-support/data"
)
model_results = df[df.model_id == "Qwen/Qwen3-0.6B"][
["neuronx_cc", "library", "library_version", "neuron_cores", "result"]
]
print(model_results)
This is just a starting point, not a full evaluation pipeline. If you build an internal report, include the model ID, software versions, result, core count, and the dataset revision (or timestamp) you used. Without that, a copied result in a ticket or spreadsheet can get separated from the setup that produced it.
Make the dataset part of the engineering record
If you're comparing multiple candidates, keep compatibility data alongside your own benchmark results. A simple record might include:
- Model ID and checkpoint revision
- Neuron compiler and library versions
- Instance type and number of cores
- Whether the model was sharded, and how
- Compilation time and steady-state measurements from your workload
- Test inputs, output length, and your quality evaluation method
The public dataset can help you decide what to test first. Your measurements should drive what you deploy. Keeping that separation avoids two common failure modes: redoing compatibility work that's already documented, or treating a public compatibility run as a substitute for workload-specific validation.
Conclusion: use compatibility data as a starting point
The Trainium model support Space helps answer a specific question: has a particular Transformers model run under a documented Neuron software stack, and what did the test actually use? With the accompanying script and public dataset, the results are more actionable than a vague compatibility label.
Just keep expectations grounded. A passing test is evidence for a bounded run, not proof of production performance. Version matching matters, compilation affects first-run time, and the stated coverage is limited by both model tasks and model size. Use the lookup to choose what to test, reproduce the setup that was tested, and then benchmark your workload before making an infrastructure decision.
Start with the model support Space, inspect the dataset and its result definitions, and treat your own measurements as the final decision point.