Models from the Transformers library can run on AWS Trainium chips thanks to TorchNeuron, a native backend for PyTorch, the open-source machine learning library. The question remains: which ones actually work? The Space torch-neuronx model support, published by Hugging Face, answers this for the most popular models on the Hub.
Search for a model
The Space opens on the Selected view, a selection of models chosen from the most downloaded and trending ones on the Hub. Each line shows the test result, the number of Trainium cores the model is distributed across, the total execution time including compilation, and the maximum memory used on the device.
To search for a specific model, simply switch to the Trending or Downloads views and enter part of its name. These views follow the Hub’s ranking, task by task, whether the model has been tested or not. When a model has been excluded, the Space specifies the reason: most often, it wasn’t a priority; sometimes, it’s genuinely incompatible. Users can also choose which versions of the Neuron compiler (neuronx-cc) and Transformers to examine.
For now, only Transformers models are tested, and the selection is limited to text generation and text-to-image to text tasks. The team plans to extend the system to diffusers, sentence transformers, and other Transformers tasks.
Run a model yourself
By selecting a model, you can see the Neuron software stack on which it was tested, the packages to install, and the run.py script used during the test. This script loads the model with from_pretrained, distributed across multiple cores if needed, then generates text with a static cache and torch.compile(backend=’neuron’).
On a trn2 instance with the same stack, for example in the AWS Neuron Deep Learning Container, you just need to install the packages and run python3 run.py. No torchrun is required; the script launches the distribution processes itself when the model is distributed. The first execution starts with compiling the model graphs. To skip this, you can copy the graphs produced during the test from the public NEFF cache bucket using the command provided by the Space, which takes the compiler version into account.
What ‘works’ means
Each test loads the model onto a trn2 instance, distributed across the smallest number of Trainium cores capable of holding its weights. Most models use tensor parallelism, which splits computations across multiple cores. Mixture-of-Experts models, without a tensor parallelism plan, only use expert parallelism. The test then generates 32 tokens, i.e., units of text, and verifies them. On a single core, these tokens are compared to those from a CPU execution of the same checkpoint. When the model is distributed, a small judge model checks that the text is fluent English.
Only models that fit on four Trainium2 devices are tested, meaning a maximum of sixteen cores. A PASS result says nothing about throughput or longer generations.
All results come from the public dataset aws-neuron/torch-neuronx-model-support, which contains a Parquet file per pair of neuronx-cc and library versions. Each row gives a model’s result, any error, the stack versions, and the run.py executed. It can be queried directly with pandas, for example for the model Qwen/Qwen3-0.6B. The dataset card specifies the exact meaning of each result.
