Artificial Intelligence

Hugging Face Launches Space to Check AWS Trainium Support

Hugging Face has introduced a dedicated tool to help developers verify model compatibility with AWS Trainium using the native PyTorch backend.

Radwan AbdulhadiEditor and technology writer

Published: October 23, 2026 · Updated: October 23, 2026 · 3 min read

Hugging Face Launches Space to Check AWS Trainium Support — Image: Hugging Face Blog

Checking Model Compatibility on AWS Trainium

Following the introduction of TorchNeuron, the native PyTorch backend for AWS Trainium via the torch-neuronx package, developers often wonder which models from the Hub can run on the hardware. To address this, Hugging Face introduced the torch-neuronx model support tool to determine model compatibility. You can explore the tool directly via the [torch-neuronx model support](https://huggingface.co/spaces/aws-neuron/torch-neuronx-model-support) space.

The interactive platform features multiple views to help users look up specific models. The default view opens on a curated selection of top models determined by download counts and trending scores. Each listed entry displays execution results, the number of Trainium cores utilized for sharding, wall time including compilation, and peak device memory consumption.

torch-neuronx model support — Image related to the report from Hugging Face Blog · Source: Hugging Face Blog

Navigating Hub Rankings and Compiler Versions

Users looking up a specific model can switch to the Trending or Downloads views, which follow standard Hugging Face Hub rankings task by task. If a model was deliberately excluded from testing, the Space provides an explanation, noting whether it was simply deprioritized or genuinely incompatible with the architecture. Developers can also filter results by selecting specific versions of the Neuron compiler and the Transformers library.

At the current stage, testing is limited to Transformers models within text-generation and image-text-to-text categories. However, future expansion plans aim to incorporate diffusers, sentence transformers, and additional task types.

Execution Scripts and Precompiled NEFF Caches

Selecting any tested model within the Space reveals a comprehensive setup guide detailing the exact Neuron stack versions, required supplementary packages, and the functional run script. The script follows the foundational method established in the initial [TorchNeuron announcement](https://huggingface.co/blog/hf-torchneuron), loading models via from_pretrained, handling necessary sharding, and executing generation tasks using a static cache alongside PyTorch compilation configurations.

When running on a Trainium2 instance—such as within the official [AWS Neuron Deep Learning Container](https://github.com/aws-neuron/deep-learning-containers#torch-neuronx)—the execution process manages distributed ranks independently without requiring separate launcher utilities like torchrun.

Because initial model runs require graph compilation, developers can bypass this compilation phase by copying prebuilt graphs directly from the public [NEFF cache](https://huggingface.co/buckets/aws-neuron/torch-neuronx-neff-cache) bucket using commands provided directly within the Space interface.

Test Methodology and Hardware Constraints

Every evaluation test loads a targeted checkpoint onto a Trainium2 instance, applying the minimum number of Trainium cores required to accommodate its weights. While standard models utilize tensor parallelism, mixture-of-experts models rely exclusively on expert parallelism when tensor-parallel plans are unavailable. Tests verify generation accuracy by producing 32 tokens, comparing single-core outputs against CPU checkpoints, and utilizing a lightweight judge model to confirm language fluency when models are sharded.

Current testing parameters cap evaluation at models fitting within four Trainium2 devices, representing a maximum of 16 cores. Successful test passes validate functional correctness rather than absolute throughput or long-generation stability.

Accessing the Underlying Dataset

All metrics and validation results displayed in the interactive tool originate from the open [aws-neuron/torch-neuronx-model-support](https://huggingface.co/datasets/aws-neuron/torch-neuronx-model-support) dataset. This repository structures data into individual Parquet files corresponding to specific combinations of compiler and library versions, enabling developers to query compatibility metrics programmatically.