TorchNeuron: running HF models on AWS Trainium is easier than ever
huggingface ·
Hugging Face has introduced TorchNeuron, a new native PyTorch backend for AWS Trainium that simplifies running models on this hardware. Unlike the previous Optimum Neuron approach which relied on lazy execution and static shapes, TorchNeuron supports eager mode, torch.compile, and distributed APIs like FSDP and TP directly within PyTorch. This integration allows Hugging Face libraries such as Transformers, Diffusers, Accelerate, PEFT, and TRL to operate on Trainium using standard interfaces without requiring custom code or dedicated trainers. For inference, text generation support covers 89% of the top 100 most-used model architectures out of the box, significantly improving upon previous capabilities. Training workflows are now streamlined, enabling users to fine-tune models with standard scripts where parallelism is handled natively by the library rather than through complex configurations. The Diffusers library also gains compatibility with tensor parallelism for image generation pipelines, including support for models like FLUX and Qwen-Image. Starting October 20, 2026, developers can install TorchNeuron via pip to access these public beta features.
huggingface