Glonce

Hugging Face launches Space showing which Hub models run on AWS Trainium

toolsThe Glonce editorsHugging Face🔥55Glonce rating55Impact40Novelty45Relevance95Surprise20Players55Credibility80
Hugging Face launches Space showing which Hub models run on AWS Trainium

Hugging Face has published a post introducing the torch-neuronx model support Space, a lookup tool for one question: does a given Hub model actually run on AWS Trainium? It builds on an earlier post about TorchNeuron, the native PyTorch backend for Trainium (the torch-neuronx package). With that backend, Transformers runs its own modeling code on Trainium as it does on any other device, so any Transformers model on the Hub may run. The Space shows which ones do, for the most popular models on the Hub.

The Space opens on a Selected view, the authors' own selection of models picked among the top Hub models 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. To look up a specific model, you 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 or not a model was tested. When a model was left out on purpose, the Space says why: most of the time it was not prioritized, and sometimes it is truly incompatible. You can also choose which versions of the Neuron compiler (neuronx-cc) and of Transformers to look at.

At the time of writing, only Transformers models are tested, and the selection is limited to text-generation and image-text-to-text models. The authors plan to expand to diffusers, sentence transformers and other transformers tasks.

For each model, the Space lists the Neuron stack it 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: it loads the model with from_pretrained, sharded when needed, then generates with a static cache and torch.compile(backend="neuron"). You run it on a Trainium2 (trn2) instance with that stack, for example in the AWS Neuron Deep Learning Container, by installing the packages and running python3 run.py. No torchrun is needed, because the script launches its own torch.distributed ranks when the model is sharded. A first run starts by compiling the model's graphs. To skip that step, you can copy the graphs the same compiler built during the test from the public NEFF cache bucket, using an hf buckets sync command that the Space supplies with the right compiler version.

The post also defines what "runs" means. Each 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. At 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.

Everything the Space shows comes from the public aws-neuron/torch-neuronx-model-support dataset, with one Parquet file per pair of neuronx-cc and library versions. Each row holds a model's result, the error when it failed, the versions of the whole stack and the run.py that ran. The post shows how to query it directly with pandas, reading the dataset from the Hub and filtering on a model_id such as Qwen/Qwen3-0.6B. The dataset card gives the exact meaning of each result.

Key facts

  • The torch-neuronx model support Space tells you whether a popular Hub model runs unmodified on AWS Trainium and provides the run.py script to reproduce it.
  • A test loads the model on a trn2 instance, generates 32 tokens and checks them: against a CPU run on a single core, or with a small judge model for fluent English when sharded.
  • Only models that fit on 4 Trainium2 devices (16 cores at most) are tested, and the selection covers text-generation and image-text-to-text models.
  • Results come from the public aws-neuron/torch-neuronx-model-support dataset, one Parquet file per pair of neuronx-cc and library versions, queryable with pandas.
  • Compiled graphs can be copied from a public NEFF cache bucket to skip the first-run compilation.

Why it matters

TorchNeuron lets Transformers run its own modeling code on Trainium, so in principle any Transformers model on the Hub may run. This Space turns that "may" into a per-model answer, with the script and the compiled-graph cache needed to reproduce it. It replaces guesswork with a lookup.

Who it affects

Practitioners who want to run Hub models on AWS Trainium, particularly those working with text-generation and image-text-to-text models, which are the only types in the selection so far. Anyone who wants the results in bulk can query the public dataset.

How to use it

Open the Space, look at the Selected view, or switch to Trending or Downloads and type part of a model's name. You can pick the neuronx-cc and Transformers versions. Select a model to get its tested Neuron stack, extra packages and run.py. On a trn2 instance with that stack, for example in the AWS Neuron Deep Learning Container, install the packages and run python3 run.py. To skip first-run compilation, sync the matching NEFF cache with the hf buckets sync command the Space gives you. For bulk queries, read the aws-neuron/torch-neuronx-model-support dataset with pandas.

How solid is it

This is a first-party Hugging Face post describing its own tool, and the test method is spelled out: 32 generated tokens, compared with a CPU run on a single core or judged for fluent English when sharded. Each dataset row records the error, the stack versions and the exact script, so results can be reproduced. The post gives no count of models tested or passed.

Risks and caveats

A PASS says nothing about throughput or about longer generations; the test checks only 32 tokens. Only models that fit on 4 Trainium2 devices (16 cores at most) are tested, and only Transformers text-generation and image-text-to-text models are in the selection. Models left out are mostly just not prioritized, though some are truly incompatible. Expansion to diffusers, sentence transformers and other tasks is planned, with no date given. The post defers the meaning of each result beyond PASS to the dataset card.