Skip to content
Neural Daily

ToolsNews

TorchNeuron makes AWS Trainium a native PyTorch device

AWS has introduced TorchNeuron, an open source native PyTorch backend for its Trainium chips. Hugging Face libraries already support it, so Trainium can be used like any other accelerator. The public beta starts on October 20.

By: Mikola Máté Patrik · · 2 min read

Hugging Face TorchNeuron cover: the yellow Hugging Face mascot hugging a Trainium chip under the TorchNeuron title
Image: Hugging Face · Press kit · Official blog thumbnail of the Hugging Face TorchNeuron post · Source

What changes

Until now, Hugging Face models ran on Trainium and Inferentia through the optimum-neuron library. It relied on the torch-neuronx package built on PyTorch XLA, which meant lazy tensors, static shapes only and a recompilation at every graph break. Optimum Neuron hid that complexity behind its own modeling classes and trainers.

TorchNeuron registers the neuron device through PyTorch's official PrivateUse1 extension point. It supports eager mode (including an experimental adaptive eager mode), torch.compile, native distributed APIs (DTensor, FSDP2, tensor parallelism) and custom kernels written with NKI. It runs on Trainium2 and Trainium3 under the Apache 2.0 license and installs with pip install torch-neuronx (PyTorch).

What it means for Hugging Face users

Transformers, Diffusers, Accelerate, PEFT and TRL now run directly on Trainium with the same classes and scripts used on GPUs. Hugging Face now recommends that Neuron users work with these libraries instead of Optimum Neuron (Hugging Face).

  • Inference: 89% of the 100 most used text generation architectures work out of the box, compared with 46% on Optimum Neuron v0.4.5. Passing device_map="neuron" is enough.
  • Training: a standard TRL SFTTrainer script runs with a few Neuron specific settings. HF tested it on 8 Trainium2 chips (32 NeuronCores, half a trn2.48xlarge), and the output is a plain Transformers checkpoint that loads on any hardware.
  • Diffusers: many pipelines run in eager and compiled mode, and FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3 also support tensor parallelism.
  • Kernels: NKI kernels are arriving on the Hub, and HF says one vision transformer kernel gives up to a 2.4x speedup on Trainium.

What to watch out for

Static shapes still matter. Generation should use a static cache, the training scripts turn off torch.compile's automatic dynamic shapes, and from TRL 1.16 batches are padded to multiples of 128 tokens. Asynchronous weight loading in Transformers does not work on Neuron yet, so it has to be disabled with HF_DEACTIVATE_ASYNC_LOAD=1.

The practical gain is clear: teams training or serving on Trainium no longer need a separate code path. Since the package is still in beta, it is worth benchmarking your own model before moving production traffic.

    Produced with the help of AI, reviewed by an editor. How we make our articles

    Reviewed and edited by the Neural Daily desk. About

    Related articles