TorchNeuron Brings Native PyTorch and Distributed Training to AWS Trainium Accelerators
Hugging Face and AWS announced TorchNeuron, a PyTorch backend that enables eager mode, torch.compile, and native distributed primitives on Trainium chips. It raises out-of-the-box support for the top 100 text-generation architectures from 46% to 89%.
Why it matters
You can now migrate existing Hugging Face training and inference workloads from NVIDIA GPUs to cheaper Trainium instances using standard classes without maintaining Optimum Neuron wrappers.
TL;DR
- 01TorchNeuron deprecates the need for Optimum Neuron by turning Trainium into a first-class PyTorch device.
- 02Out-of-the-box support for the top 100 text generation models jumps from 46% to 89%.
- 03Fine-tuning scripts using TRL SFTTrainer produce standard checkpoints loadable on any GPU.
- 04Dynamic shapes and async weight loading must be disabled during current beta training.
Key facts
- Public Beta Availability
- October 20, 2026
- Top 100 Text-Gen Architecture Support
- 89% (up from 46% with Optimum Neuron)
- Tested Cluster Size
- 8 Trainium2 chips (32 NeuronCores, half trn2.48xlarge)
- Installation Target
- pip install torch-neuronx
Moving from Optimum Neuron to Native PyTorch
Until now, running Hugging Face models on AWS Trainium required optimum-neuron, which implemented custom modeling abstractions around PyTorch XLA. TorchNeuron integrates directly into core PyTorch as an acceleration backend. This exposes neuron as a native device supporting eager execution, torch.compile, and standard PyTorch distributed APIs like FSDP2 and DTensor.
The shift expands model coverage significantly: 89% of the top 100 text-generation architectures now run out-of-the-box, compared to 46% under optimum-neuron v0.4.5. Popular vision and diffusion models—including FLUX.1-dev, FLUX.2, and Qwen-Image-2.1—can also run with native tensor parallelism directly inside Diffusers.
Distributed Training via SFTTrainer
Fine-tuning with TRL requires no Neuron-specific forks. Parallelism is defined inside model_init_kwargs through DistributedConfig(tp_size=..., fsdp_size=...). The model shards automatically across tensor-parallel and FSDP dimensions upon loading weights:
- Static shapes: Automatic dynamic shapes in Dynamo must be turned off via
torch._dynamo.config.automatic_dynamic_shapes = False. Inputs are padded to multiples of 128 tokens in TRL ≥ 1.16 to reuse compiled graphs. - Single-flag compilation: Running with
--torch_compile trueautomatically targets the Neuron backend. - Deactivate async load: Set
HF_DEACTIVATE_ASYNC_LOAD=1, as asynchronous weight streaming is currently unsupported.
Launching training across 8 Trainium2 chips (32 NeuronCores on a trn2.48xlarge instance) uses torchrun with --nproc_per_node=32 alongside TP_SIZE=4 and FSDP_SIZE=8. Checkpoints produced by trainer.save_model remain vanilla Hugging Face artifacts that can be reloaded onto GPUs without format conversion.
Try it in 2 minutes
pip install torch-neuronx
# Distributed fine-tuning launch on 8 Trainium2 chips (32 cores)
TP_SIZE=4 FSDP_SIZE=8 HF_DEACTIVATE_ASYNC_LOAD=1 uv run torchrun --nproc_per_node=32 sft_neuron_example.pybash
✓ When to use
- Fine-tuning large language models on AWS Trainium using standard Hugging Face TRL scripts.
- Serving open-weight text generation models with tensor parallelism on AWS trn2 instances.
- Decoupling PyTorch training pipelines from proprietary NVIDIA CUDA APIs.
✕ When NOT to use
- Workloads requiring automatic dynamic shapes without input token sequence padding.
- Inference pipelines that depend on Transformers asynchronous weight streaming.
- Workstations and cloud providers without dedicated AWS Trainium or Inferentia silicon.
What to do today
- Install the beta backend via pip install torch-neuronx on AWS Trainium EC2 instances.
- Update fine-tuning scripts to replace Optimum Neuron classes with native transformers and TRL SFTTrainer.
- Set HF_DEACTIVATE_ASYNC_LOAD=1 in your distributed launch environment.