Does Your Model Run on AWS Trainium? Look It Up
Hugging Face launches a tool to check which Transformers models run unmodified on AWS Trainium, providing run scripts and performance metrics via a public Space.
Hugging Face has introduced a Space that tracks which Transformers models from its Hub can run natively on AWS Trainium using the torch-neuronx backend. The tool displays whether a model runs without modification, the number of Trainium cores required, compilation and runtime duration, and peak device memory usage. Users can search models by name in the Trending or Downloads views, which reflect the Hub’s rankings. The Space also explains why certain models were excluded, such as lack of prioritization or incompatibility.
The Space provides detailed instructions for running each model, including the Neuron stack version, required packages, and a run.py script that follows Hugging Face’s TorchNeuron recipe. The script loads the model with from_pretrained, applies sharding when necessary, and generates text using a static cache and torch.compile with the Neuron backend. Users can skip compilation by downloading pre-built graphs from the public NEFF cache bucket, with commands provided for the correct compiler version.
Testing is performed on Trainium2 (trn2) instances, with models sharded across the minimum number of Trainium cores needed to fit their weights. Most models use tensor parallelism, while mixture-of-experts models without a tensor-parallel plan rely solely on expert parallelism. The Space validates model outputs by generating 32 tokens and comparing them to CPU runs for single-core tests or using a small judge model for sharded tests to ensure fluent English output.
The data behind the Space is sourced from the public aws-neuron/torch-neuronx-model-support dataset, which contains Parquet files for each pair of neuronx-cc and library versions. Each file includes a model’s result, error details if it failed, stack versions, and the run.py script used. Users can query the dataset directly, and the dataset card clarifies the meaning of each result. Currently, only text-generation and image-text-to-text models are tested, with plans to expand to diffusers, sentence transformers, and other tasks.