AWS Trainium Model Checker Reveals Which Transformers Models Run

Summary

Published October 23, 2026, the torch-neuronx model support Space shows whether selected popular Hugging Face Hub Transformers models run unmodified on AWS Trainium and provides reproducible scripts. Its Selected view draws from download and trending leaders; searchable Trending and Downloads views retain task-level Hub rankings, include untested models, explain deliberate omissions, and filter by neuronx-cc and Transformers versions. Rows show results, Trainium core counts, compilation-inclusive wall time and peak device memory. Coverage currently includes text-generation and image-text-to-text models, with diffusers, sentence transformers and other Transformers tasks planned.

Each entry identifies the tested Neuron stack, extra packages and run.py. On a Trainium2, or trn2, instance such as the AWS Neuron Deep Learning Container, the script loads with from_pretrained, shards when necessary, generates through a static cache and torch.compile with the neuron backend, and launches its own torch.distributed ranks without torchrun. First runs compile graphs, but matching graphs built with the same compiler version can be synced from the public NEFF cache.

Tests load each model onto the fewest cores holding its weights, using tensor parallelism or expert parallelism for mixture-of-experts models lacking a tensor-parallel plan, then generate and check 32 tokens. Single-core output is compared with a CPU run; sharded text is assessed for fluent English by a small judge model. Testing is capped at four Trainium2 devices, or 16 cores. A PASS establishes neither throughput nor longer-generation quality. The public aws-neuron/torch-neuronx-model-support dataset stores one Parquet file per neuronx-cc and library-version pair, including results, errors, full-stack versions and scripts, with result definitions in its dataset card.