The torch-neuronx model support Space on the Hugging Face Hub is a lookup tool for whether Transformers models run unmodified on AWS Trainium. It follows the earlier TorchNeuron announcement, which introduced torch-neuronx as the native PyTorch backend for Trainium. Because Transformers runs its own modeling code on any device, any model on the Hub may run, but the Space verifies which popular ones actually pass.

The Space opens on a Selected view built from top models by downloads and trending score. Each row shows the result, the number of Trainium cores the model is sharded over, wall time including compilation, and peak device memory. Users can switch to Trending or Downloads views to follow the Hub’s own rankings and search by model name. The Space also lets you choose the versions of neuronx-cc and Transformers to inspect. When a model was intentionally left out, it explains why; usually it was simply not prioritized, occasionally it is genuinely incompatible. At the time of writing, only Transformers text-generation and image-text-to-text models are covered, with plans to add diffusers, sentence transformers, and other tasks.

Selecting a model gives the tested Neuron stack, the packages to install, and the run.py script that was used. The script loads the model with from_pretrained, shards if necessary, and generates using a static cache with torch.compile(backend=’neuron’). On a Trainium2 (trn2) instance with the same stack, you can install the listed packages and run python3 run.py directly; no torchrun is required because the script launches its own torch.distributed ranks when the model is sharded. A first run compiles the model’s graphs. To skip that step, copy the graphs from the public NEFF cache bucket using the provided command with the matching compiler version, for example hf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc- /tmp/neff_cache.

The test procedure loads each model on a trn2 instance, sharded over the smallest number of cores that can hold its weights. Most models use tensor parallelism; mixture-of-experts models without a tensor-parallel plan use expert parallelism only. The test generates 32 tokens and checks them. For a single core, tokens are compared against a CPU run of the same checkpoint; for sharded models, a small judge model checks that the text is fluent English. Only models that fit on four Trainium2 devices, or 16 cores, are tested, and a PASS says nothing about throughput or longer generations.

All displayed results come from the public aws-neuron/torch-neuronx-model-support dataset, stored as one Parquet file per pair of neuronx-cc and library versions. Rows contain the model ID, result, error if it failed, versions of the whole stack, and the run.py used. The dataset can be queried directly with pandas, e.g., df = pd.read_parquet(‘hf://datasets/aws-neuron/torch-neuronx-model-support/data’) followed by a filter on model_id. The dataset card documents the exact meaning of each result.