{"slug": "does-your-model-run-on-aws-trainium-look-it-up", "title": "Does Your Model Run on AWS Trainium? Look It Up", "summary": "Hugging Face and AWS Neuron published a torch-neuronx model support Space and public dataset that records which Transformers models run unmodified on AWS Trainium, tested on trn2 instances with up to 16 cores. Each entry reports the number of Trainium cores a model is sharded over, wall time including compilation, and peak device memory, and links a run.py script using from_pretrained, a static cache and torch.compile(backend=\"neuron\"). Testing is currently limited to text-generation and image-text-to-text Transformers models that fit on 4 Trainium2 devices, and a PASS says nothing about throughput or longer generations.", "body_md": "[Updated   •  60](https://huggingface.co/datasets/aws-neuron/torch-neuronx-model-support)    \n\n# \n\t\tDoes Your Model Run on AWS Trainium? Look It Up\n\t\n\n [Update on GitHub](https://github.com/huggingface/blog/blob/main/torch-neuronx-model-support.md)\n\nIn a [previous post](https://huggingface.co/blog/hf-torchneuron), we introduced TorchNeuron, the native PyTorch backend for AWS Trainium (the `torch-neuronx` package). With it, Transformers runs its own modeling code on Trainium, as it does on any other device, so any Transformers model on the Hub may run. Which ones actually do?\n\nThe [torch-neuronx model support](https://huggingface.co/spaces/aws-neuron/torch-neuronx-model-support) Space answers that question for the most popular models on the Hub. It tells you whether a model runs unmodified on Trainium, and gives you the script to run it yourself.\n\n## \n\t\tLook up a model\n\t\n\nThe Space opens on the **Selected** view: our own selection of models picked among the top models on the Hub by downloads and by trending score. Each row gives the result, the number of Trainium cores the model is sharded over, the wall time of the run (compilation included) and the peak device memory.\n\nTo look up a given model, switch to the **Trending** or **Downloads** view and type part of its name. These views follow the Hub's own rankings, one task at a time, whether a model was tested or not. When a model was left out on purpose, the Space tells you why (most of the time it was not prioritized, sometimes it is truly incompatible). You can also pick the versions of the Neuron compiler (`neuronx-cc`) and of Transformers to look at.\n\nAt 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.\n\n## \n\t\tRun it yourself\n\t\n\nSelect a model to see how to run it. The Space lists the Neuron stack the model was tested on, the packages to install on top of it, and the `run.py` script that ran. The script follows the recipe of the [TorchNeuron announcement](https://huggingface.co/blog/hf-torchneuron): it loads the model with `from_pretrained`, sharded when it needs to be, then generates with a static cache and `torch.compile(backend=\"neuron\")`.\n\nOn a Trainium2 (`trn2`) instance with that stack, for example in the [AWS Neuron Deep Learning Container](https://github.com/aws-neuron/deep-learning-containers#torch-neuronx), install the packages, then run:\n\n```\npython3 run.py\n```\n\nNo need for `torchrun`: the script launches its own `torch.distributed` ranks when the model is sharded.\n\nA first run starts by compiling the model's graphs. To skip that step, copy the graphs the same compiler built during the test, from the public [NEFF cache](https://huggingface.co/buckets/aws-neuron/torch-neuronx-neff-cache) bucket. The Space gives you the command, with the right compiler version:\n\n```\nhf buckets sync hf://buckets/aws-neuron/torch-neuronx-neff-cache/neuronxcc-<version> /tmp/neff_cache\n```\n\n## \n\t\tWhat \"runs\" means\n\t\n\nEach 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 without a tensor-parallel plan use expert parallelism only. The test then generates 32 tokens and checks them. On a single core, the 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 text is fluent English.\n\nAt the time of writing, only models that fit on 4 Trainium2 devices are tested, which means 16 cores at most. A `PASS` says nothing about throughput, or about longer generations.\n\n## \n\t\tThe results are a dataset\n\t\n\nEverything the Space shows comes from the public [`aws-neuron/torch-neuronx-model-support`](https://huggingface.co/datasets/aws-neuron/torch-neuronx-model-support) dataset, with one Parquet file per pair of `neuronx-cc` and library versions.\n\nEach row holds a model's result, the error when it failed, the versions of the whole stack and the `run.py` that ran. You can query it directly:\n\n``` python\nimport pandas as pd\n\ndf = pd.read_parquet(\"hf://datasets/aws-neuron/torch-neuronx-model-support/data\")\ndf[df.model_id == \"Qwen/Qwen3-0.6B\"][[\"neuronx_cc\", \"library\", \"library_version\", \"neuron_cores\", \"result\"]]\n```\n\nThe [dataset card](https://huggingface.co/datasets/aws-neuron/torch-neuronx-model-support) gives the exact meaning of each result.\n\nNext time you wonder whether a model runs on Trainium, [look it up](https://huggingface.co/spaces/aws-neuron/torch-neuronx-model-support).", "url": "https://wpnews.pro/news/does-your-model-run-on-aws-trainium-look-it-up", "canonical_source": "https://huggingface.co/blog/torch-neuronx-model-support", "published_at": "2026-10-23 00:00:00+00:00", "updated_at": "2026-10-23 13:42:07.504139+00:00", "lang": "en", "topics": ["ai-infrastructure", "ai-chips", "large-language-models", "mlops", "ai-tools"], "entities": ["Hugging Face", "AWS Trainium", "torch-neuronx", "Trainium2", "neuronx-cc", "Transformers", "AWS Neuron", "TorchNeuron"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/does-your-model-run-on-aws-trainium-look-it-up", "markdown": "https://wpnews.pro/news/does-your-model-run-on-aws-trainium-look-it-up.md", "text": "https://wpnews.pro/news/does-your-model-run-on-aws-trainium-look-it-up.txt", "jsonld": "https://wpnews.pro/news/does-your-model-run-on-aws-trainium-look-it-up.jsonld"}}