AWS Trainium Model Support: Hugging Face Launches a torch-neuronx Compatibility Lookup

TL;DR — Hugging Face and AWS have launched the torch-neuronx model support Space, a public lookup that tells you whether a given Transformers model on the Hugging Face Hub runs unmodified on AWS Trainium, and hands you the exact script and stack that were used to test it. It turns "does this model work on Trainium?" from a trial-and-error question into a queryable dataset.

The Space builds on TorchNeuron, the native PyTorch backend for AWS Trainium shipped in the torch-neuronx package, which was introduced in a previous Hugging Face post. With TorchNeuron, Transformers runs its own modeling code on Trainium as it does on any other device — so in principle any Transformers model on the Hub may run. The new Space answers which ones actually do, for the most popular models on the Hub.

What was announced

The aws-neuron/torch-neuronx-model-support Space is a public, browsable compatibility table for AWS Trainium. For each model listed, it reports:

  • Whether the model runs unmodified on Trainium (a PASS result)
  • The number of Trainium cores the model was sharded over
  • The wall time of the run, compilation included
  • The peak device memory used

For every model, the Space also exposes the script that was run, so the result can be reproduced rather than trusted.

How to look up a model

The Space opens on a Selected view, described as the authors' own selection of models picked among the top models on the Hub by downloads and by trending score.

The Trending and Downloads views follow the Hub's own rankings, one task at a time. These views list models whether they were tested or not, so a missing entry is never ambiguous: when a model was deliberately left out, the Space states why.

At the time of writing, only Transformers models are tested, and our selection is limited to text-generation and image-text-to-text models. We plan to expand to diffusers, sentence transformers, and other transformers tasks.

The Space also lets you select the versions of the Neuron compiler (neuronx-cc) and of Transformers you want to inspect, so results can be filtered by stack version rather than read as a single undated snapshot.

To find a specific model, switch to the Trending or Downloads view and type part of its name. The space reports why a model was excluded — most often because it was not prioritized, and sometimes because it is genuinely incompatible with Trainium.

How to run the model yourself

Selecting a model in the Space reveals how to run it. The Space lists the Neuron stack the model was tested on, the packages to install on top of that stack, and the run.py script that actually ran.

The script follows the recipe from the TorchNeuron announcement: it loads the model with from_pretrained, sharded when the model requires it, then generates with a static cache and torch.compile(backend="neuron").

On a Trainium2 (trn2) instance carrying that stack — for example, inside the AWS Neuron Deep Learning Container — you install the listed packages and then run:

python3 run.py

No torchrun invocation is needed. When the model is sharded, the script launches its own torch.distributed ranks.

Skipping compilation with the NEFF cache

A first run begins by compiling the model's graphs, which is where the reported wall time is largely spent. To skip that step, you can copy the graphs that the same compiler built during the test from the public NEFF cache bucket. The Space supplies the command with the correct compiler version filled in:

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

Using the cache means a reproduction can skip recompilation as long as it uses the same neuronx-cc version as the original test.

What "runs" means

The compatibility claim is narrow and specific, and the Space documents it precisely.

Each test loads the model on a trn2 instance, sharded over the smallest number of Trainium cores that can hold its weights. Most models are sharded with tensor parallelism; mixture-of-experts models that lack a tensor-parallel plan use expert parallelism only. The test then generates 32 tokens and checks them.

Correctness checking depends on the sharding configuration:

  • On a single core, the generated tokens are compared with those of a CPU run of the same checkpoint.
  • When the model is sharded, a small judge model checks that the generated text is fluent English.

The scope of a PASS is bounded deliberately. At the time of writing, only models that fit on 4 Trainium2 devices are tested, which means at most 16 cores. A PASS says nothing about throughput, and nothing about longer generations.

The results are a dataset

Everything the Space displays comes from the public aws-neuron/torch-neuronx-model-support dataset, which stores one Parquet file per pair of neuronx-cc and library versions. Each row holds a model's result, the error when it failed, the versions of the whole stack, and the run.py that ran.

Because it is a dataset rather than a rendered table, the results 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 exact meaning of each result value, so a PASS or failure can be interpreted against the precise methodology rather than inferred.

Coverage limits and what comes next

Four constraints bound what the Space currently proves:

  1. Only Transformers models are tested.
  2. Only text-generation and image-text-to-text tasks are covered.
  3. Only models that fit on 4 Trainium2 devices — 16 cores — are tested.
  4. A PASS reports correctness for 32 generated tokens, not throughput or long-context behaviour.

The stated plan is to expand coverage to diffusers, sentence transformers, and other Transformers tasks.

Why this matters

Hardware portability questions are usually answered by tribal knowledge, ad-hoc issues, or expensive trial runs. Publishing Trainium compatibility as a versioned, queryable dataset with the reproduction script attached changes the economics of that question: the answer is a lookup, and the evidence is inspectable.

Two design choices make the resource more useful than a leaderboard would be. First, results are keyed to specific neuronx-cc and library versions, so compatibility can be evaluated against the stack you actually plan to deploy. Second, the dataset records failures and the errors that produced them, alongside successes, which makes previously undocumented incompatibilities searchable.

The practical takeaway: before running a model on AWS Trainium, check the aws-neuron/torch-neuronx-model-support Space — or query its Parquet dataset directly — and reuse the published run.py and NEFF cache entry to reproduce the result on a trn2 instance.

Sources