Hugging Face launches tool to check AWS Trainium model compatibility

Hugging Face has introduced a new Space to help developers determine if popular models from the Hub are compatible with AWS Trainium via the torch-neuronx package. The tool provides details on model sharding, core usage, memory consumption, and execution time. It also provides the specific scripts and environment configurations required to run tested models on Trainium2 instances.

Testing and verification details

• The tool tracks model sharding across Trainium cores, wall time (including compilation), and peak device memory.
• Compatibility testing currently focuses on text-generation and image-text-to-text Transformers models.
• Successful runs use from_pretrained, sharding where necessary, and torch.compile(backend="neuron") with a static cache.
• Results are exposed via a public Parquet dataset (aws-neuron/torch-neuronx-model-support) for direct querying.
• Testing is limited to models fitting on up to 4 Trainium2 devices (16 cores).

Why it matters

This tool lowers the barrier for deploying open-source AI models on specialized AWS hardware by providing clear compatibility data and ready-to-use execution scripts.