Hugging Face announces native support for AWS’s Trainium chip via TorchNeuron, a new backend for PyTorch. Transformers, Diffusers, and Kernels can now run directly on this chip, which should simplify running models on Trainium and sharing them on the Hugging Face Hub. According to the source, a public beta will be available starting October 20, 2026.
TorchNeuron is an extension of PyTorch, the reference library for machine learning model development. It adds a backend dedicated to Trainium to PyTorch, with the “eager” execution mode, torch.compile (a code optimization tool), and native distributed APIs like FSDP and TP. Researchers and developers will be able to install it with the command pip install torch-neuronx.
A Shift in Approach for Hugging Face Libraries
Until now, using Trainium and Inferentia, two chips designed by AWS—one for cost-effective model training and the other for large-scale inference—required the optimum-neuron library. This library relied on the torch-neuronx package, itself based on PyTorch XLA, which forced users to think “like a compiler”: lazy tensors, static shapes only, and recompilation at every graph break. Hugging Face states that optimum-neuron was used to hide this complexity via specifically maintained code.
TorchNeuron changes this approach by integrating directly into PyTorch, which sees Neuron as a native device. Trainium then becomes an accelerator like any other for Transformers, Diffusers, Accelerate, PEFT, and TRL, with the same classes and scripts. Hugging Face now recommends Neuron users to go through these libraries rather than Optimum Neuron.
For text generation, support is broad: according to Hugging Face, 89% of the top 100 most-used text generation model architectures are compatible out of the box, compared to 46% with Optimum Neuron v0.4.5. Simply passing the argument device_map=neuron moves the model to Neuron devices.
Training, Diffusion, and Compute Kernels
For training, a standard fine-tuning script with TRL works on a Trainium instance with a few Neuron-specific adjustments. Hugging Face specifies that the published example uses Qwen/Qwen3-8B and the HuggingFaceH4/ultrachat_200k dataset, and was run on 8 Trainium2 chips, i.e., 32 NeuronCores, which is half of a trn2.48xlarge instance. The produced checkpoint is a standard Transformers checkpoint, usable on other hardware.
Several Diffusers pipelines are compatible with TorchNeuron, in both “eager” and compiled modes. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1, and H3 also support tensor parallelism, which distributes each model layer across multiple cores. Hugging Face indicates that other pipelines will gradually become compatible and that a guide allows users to add this support themselves.
Finally, the Kernels library will host compute kernels optimized for Neuron, written with the Neuron Kernel Interface (NKI). A kernel for vision transformers, published on the Hub, reportedly enables acceleration “up to 2.4x” on Trainium. Hugging Face thanks the AWS team for their work on TorchNeuron and their responsiveness to community feedback.
