Hugging Face launches Trainium2 model compatibility lookup

Hugging Face has launched a public lookup for selected Transformers models on AWS Trainium2 that shows whether each tested checkpoint completes a defined compatibility run and supplies the script used to reproduce it. The tool turns a general framework-support claim into a model-level record tied to a specific software stack.

Each result records the outcome, the number of Trainium cores used to shard the model, total wall time including compilation, and peak device memory. Users can filter results by Neuron compiler and Transformers versions. A model’s detail page lists the tested Neuron stack, extra packages and the run.py script that executed. The public dataset keeps failed runs and their error records instead of omitting them.

The test uses a narrow definition of “runs.” Hugging Face loads each checkpoint on a Trainium2 instance, shards it across the smallest number of Trainium cores able to hold its weights, and generates 32 tokens with a static cache and torch.compile using the Neuron backend. Most models use tensor parallelism; mixture-of-experts models without a tensor-parallel plan use expert parallelism only.

Validation depends on the configuration. If a model fits on one core, the generated tokens are compared with a CPU run of the same checkpoint. For a sharded model, a small judge model checks whether the output is fluent English instead of comparing it token for token with a CPU result. A PASS therefore establishes only that the model completed the disclosed compatibility check under the recorded stack. Hugging Face says it does not establish throughput or behavior over longer generations. The wall-time and memory figures document that run; they are not production latency, throughput or cost benchmarks.

The lookup also provides a reproduction path. Hugging Face says users can run the published script on a Trainium2 instance with the matching stack. To avoid compilation on the first run, they can copy graphs built by the same compiler from a public NEFF cache. AWS’s initial TorchNeuron release notes require static shapes for torch.compile, say a new input shape may trigger another compilation, and note that first execution includes compilation when no matching cached artifact exists.

The public aws-neuron/torch-neuronx-model-support dataset backs the Space. Hugging Face organizes the results as one Parquet file for each pairing of neuronx-cc and library versions. Each row records a model result, any error from a failed run, full software-stack versions and the script used, making passes and failures traceable to a reproducible configuration.

The lookup builds on Hugging Face’s TorchNeuron integration, which made Trainium a native neuron device for Hugging Face libraries and recommended using those libraries directly instead of the earlier Optimum Neuron abstraction. DataPhoenix previously covered the native TorchNeuron path for AWS Trainium. AWS launched TorchNeuron in public beta as an open-source native PyTorch backend for Amazon EC2 Trainium2 and Trainium3 instances.

Coverage remains limited at launch. The lookup tests only selected Transformers models in the text-generation and image-text-to-text tasks; Diffusers, Sentence Transformers and other Transformers tasks are outside the published results. Hugging Face also limits testing to models that fit on no more than four Trainium2 devices, or 16 Trainium cores, and has not published a frozen aggregate pass rate.