TorchNeuron is a new PyTorch plugin that provides a native backend for AWS Trainium (Neuron), exposing Trainium as a PyTorch device with eager execution, torch.compile support and native distributed APIs like DTensor and FSDP2. Hugging Face libraries such as Transformers, Diffusers, Accelerate, PEFT and TRL can now run on Trainium directly, without routing through Optimum Neuron.
Availability
TorchNeuron will be offered as a public beta starting October 20, 2026. It will be available for installation via pip install torch-neuronx.
What changes compared to the previous approach?
Until now, Hugging Face support for AWS Trainium and Inferentia relied on Optimum Neuron, which built on the torch-neuronx package based on PyTorch XLA. That approach required a compiler-like mindset (lazy tensors, static shapes, recompilation on graph breaks) and Optimum Neuron hid that complexity with custom modeling and trainer code.
TorchNeuron integrates directly into PyTorch and exposes the neuron device natively: eager execution, PyTorch’s distributed primitives (DTensor, FSDP2) and torch.compile compatibility. For Hugging Face libraries, Trainium behaves like any other device — Transformers, Diffusers, Accelerate, PEFT and TRL run on it with the same classes and scripts used for other accelerators. Hugging Face now recommends using these libraries directly rather than going through Optimum Neuron.
Transformers — inference and text generation
Transformers now supports text generation natively for a wide range of models on TorchNeuron. According to the announcement, 89% of the top 100 most-used text-generation model architectures are supported out of the box with TorchNeuron, compared to 46% with Optimum Neuron v0.4.5. Using the device_map=neuron argument moves the model onto Neuron devices.
The post includes example code demonstrating tensor parallelism and the use of a static cache for generation; static caching is recommended on Neuron to avoid frequent recompilations.
Training — fine-tuning on Trainium
Because Neuron is presented as a native PyTorch device, training on Trainium uses the same Hugging Face libraries as any other accelerator: Transformers, Accelerate, PEFT and TRL. There is no need for a dedicated Neuron trainer or a fork of model code; a standard TRL SFTTrainer script runs on Trainium with a few Neuron-specific settings.
Key points to know when training on Trainium:
- Static shapes: the example disables torch.compile automatic dynamic shapes and uses wrapped packing for sequences up to 2048 tokens. With TRL ≥ 1.16, batches on Neuron are padded to a multiple of 128 tokens, reducing the number of input shapes and enabling reuse of compiled graphs.
- torch.compile: passing --torch_compile true is sufficient; torch.compile detects the device and selects the Neuron backend automatically.
- Async weight loading: Transformers’ built-in asynchronous weight loading is not supported on Neuron, so HF_DEACTIVATE_ASYNC_LOAD=1 must be used to disable it.
The article shows a run command where TP_SIZE=4 and FSDP_SIZE=8 create 32 processes (4 x 8 = 32), matching --nproc_per_node=32. The authors ran on 8 Trainium2 chips (32 NeuronCores), which corresponds to half of a trn2.48xlarge instance. Training completes and trainer.save_model writes a standard Transformers checkpoint that is not Neuron-specific and can be pushed to the Hugging Face Hub and loaded on other hardware.
Diffusers — image generation
Many Diffusers pipelines are compatible with TorchNeuron in both eager and compiled modes. Models named in the announcement — FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3 — support tensor parallelism to leverage Neuron devices. The article includes example code illustrating TP=8 DeviceMesh initialization, transformer parallel configuration, compile_repeated_blocks with the Neuron backend and image generation.
Example run: NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py. The announcement notes more pipelines will become TP-compatible and points to a guide for adding support if a given pipeline is not yet compatible.
Kernels and custom Neuron operations
The Neuron team’s kernels are being published on the Hub to accelerate models. The Neuron Kernel Interface (NKI) allows developers to implement custom kernels and share them on the Hub. The announcement cites an example kernel that can accelerate vision transformers on Trainium by up to 2.4x.
Acknowledgments
The post thanks the AWS team for their work on TorchNeuron and for responding to Hugging Face feedback to improve the integration.
Conclusion
TorchNeuron removes the previous compiler-like layer and gives a native PyTorch experience on Trainium: eager execution, torch.compile support and direct use of standard Hugging Face libraries for inference and training. With the public beta starting October 20, 2026, researchers and developers will be able to install TorchNeuron and more easily run and share Trainium models via the Hugging Face Hub.
