TorchNeuron: running HF models on AWS Trainium is easier than ever

Hugging Face and AWS will release TorchNeuron, a native PyTorch backend for AWS Trainium, as a public beta starting Oct. 20, 2026, installable via `pip install torch-neuronx`. With TorchNeuron, 89% of the top 100 most-used text-generation model architectures are supported out of the box in Transformers, up from 46% with Optimum Neuron v0.4.5, and Hugging Face now recommends Neuron users run Transformers, Diffusers, Accelerate, PEFT and TRL directly instead of going through Optimum Neuron. The backend exposes `neuron` as a native device with eager execution, `torch.compile` and PyTorch distributed primitives (DTensor, FSDP2), replacing the PyTorch XLA-based lazy-tensor approach that required static shapes and recompilation for every graph break.

TorchNeuron: running HF models on AWS Trainium is easier than ever Update on GitHub https://github.com/huggingface/blog/blob/main/hf-torchneuron.md TorchNeuron TorchNeuron https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/index.html is a PyTorch plugin that brings a new backend for Trainium enabling support for eager mode, torch.compile and native distributed APIs like FSDP and TP. Read the Pytorch blog https://pytorch.org/blog/introducing-torchneuron-native-pytorch-backend-for-aws-trainium/ to learn more about TorchNeuron. Starting Oct. 20, 2026, TorchNeuron will be available as a public beta. Researchers and developers will be able to install it with pip install torch-neuronx . New backend, native interface in HF libraries To date, the Hugging Face experience on AWS Trainium and Inferentia — AI chips designed respectively for cost-effective training and inference at scale — has been delivered through optimum-neuron https://github.com/huggingface/optimum-neuron . This library relies on the torch-neuronx package, which is based on PyTorch XLA. It was the natural choice at the time, but it meant thinking like a compiler: lazy tensors instead of eager execution, static shapes only, a recompilation for every graph break. Optimum Neuron existed to hide that complexity, and maintained custom code modeling classes, trainers, etc to do so. TorchNeuron changes this. It integrates directly into PyTorch and exposes neuron as a native device, with eager execution, PyTorch's own distributed primitives DTensor, FSDP2 and torch.compile . For the Hugging Face libraries, Trainium becomes a device like any other: Transformers, Diffusers, Accelerate, PEFT and TRL run on it directly, with the same classes and scripts you would use on any other accelerator. From now on, we recommend that Neuron users use these libraries directly instead of going through Optimum Neuron. Transformers Inference Text generation is now natively supported in Transformers for a wide range of models. With TorchNeuron, 89% of the top 100 most-used text-generation model architectures are supported out of the box, compared to 46% with Optimum Neuron v0.4.5. You just need to pass the device map=neuron argument in your code and your model will be moved to the Neuron devices. The following is an example of text-generation with tensor parallelism, which will enable to fully benefit from your Neuron instance. Note that it uses a static cache for generation as this is strongly recommended with Neuron to prevent frequent recompilations. python import os import torch import torch.distributed as dist import torch neuronx noqa: F401 from torch.distributed import DeviceMesh from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig from transformers.distributed import DistributedConfig from transformers.generation.configuration utils import CompileConfig MODEL ID = "Qwen/Qwen3.8-27B" MAX SEQ LEN = 1024 PREFILL CHUNK SIZE = 512 dist.init process group backend="neuron" world size = dist.get world size model = AutoModelForCausalLM.from pretrained MODEL ID, dtype=torch.bfloat16, distributed config=DistributedConfig tp size=world size , device mesh=DeviceMesh "neuron", list range world size , .eval tokenizer = AutoTokenizer.from pretrained MODEL ID if tokenizer.pad token is None: tokenizer.pad token = tokenizer.eos token tokenizer.padding side = "left" inputs = tokenizer.apply chat template {"role": "user", "content": "The future of AI is"} , add generation prompt=True, padding=True, tokenizer kwargs={"pad to multiple of": PREFILL CHUNK SIZE}, return tensors="pt", generation config = GenerationConfig max new tokens=128, do sample=False, use cache=True, pad token id=tokenizer.pad token id, eos token id=tokenizer.eos token id, cache implementation="static", max cache len=MAX SEQ LEN, prefill chunk size=PREFILL CHUNK SIZE, compile config=CompileConfig backend="neuron", fullgraph=False, dynamic=False , with torch.no grad : output ids = model.generate inputs "input ids" , attention mask=inputs "attention mask" , generation config=generation config, if dist.get rank == 0: prompt len = inputs "input ids" .shape 1 print tokenizer.decode output ids 0, prompt len: , skip special tokens=True Training Now that Neuron is a native PyTorch device, training on Trainium uses the same Hugging Face libraries as any other accelerator: Transformers, Accelerate, PEFT and TRL. You no longer need a dedicated trainer or a Neuron-specific fork of the model code. A standard TRL fine-tuning script runs on a Trainium instance with only a few Neuron-specific settings. Here is a supervised fine-tuning script built on TRL's SFTTrainer . If you have fine-tuned a model with TRL on GPUs before, it should look familiar. python import os import torch from datasets import load dataset from transformers.distributed import DistributedConfig from trl import ModelConfig, ScriptArguments, SFTConfig, SFTTrainer, TrlParser def main script args, training args, model args : The Neuron backend can't compile Dynamo's automatic dynamic shapes yet https://github.com/huggingface/trl/issues/7660 torch. dynamo.config.automatic dynamic shapes = False training args.model init kwargs = { "revision": model args.model revision, "trust remote code": training args.trust remote code, "attn implementation": model args.attn implementation, "dtype": model args.dtype, "device map": None, } if int os.environ.get "WORLD SIZE", "1" 1: The model is sharded on the TP x FSDP grid as its weights are loaded training args.model init kwargs "distributed config" = DistributedConfig tp size=int os.environ.get "TP SIZE", "1" , fsdp size=int os.environ.get "FSDP SIZE", "1" , Prepare dataset train dataset = load dataset script args.dataset name, name=script args.dataset config, split=script args.dataset train split, UltraChat also has a string prompt column, which makes TRL treat it as a prompt-completion dataset. train dataset = train dataset.select columns "messages" Training trainer = SFTTrainer model=model args.model name or path, args=training args, train dataset=train dataset, trainer.train trainer.save model training args.output dir trainer.end if name == " main ": parser = TrlParser ScriptArguments, SFTConfig, ModelConfig script args, training args, model args = parser.parse args and config main script args, training args, model args Parallelism is handled natively by Transformers. You put a DistributedConfig with the tensor-parallel TP and FSDP sizes in model init kwargs , SFTTrainer passes it to from pretrained , and the model is sharded as its weights are loaded. The Trainer sees that the model is already distributed and does not wrap it again. You don't need an --fsdp flag or an Accelerate FSDP config. Nothing in the script is tied to Trainium: the same code runs on any hardware. A few things are worth knowing when you run it on Trainium: - Static shapes. The script disables automatic dynamic shapes in torch.compile and uses wrapped packing to produce sequences of up to 2048 tokens. With TRL ≥ 1.16, batches on Neuron are padded to a multiple of 128 tokens, limiting the number of input shapes. Compiled graphs are reused for subsequent steps with the same shapes. - torch.compile on. Passing --torch compile true is all it takes: on a Trainium instance, torch.compile detects the device and picks the Neuron backend automatically, with nothing else to configure. - Async weight loading off. The built-in asynchronous weight loading in Transformers is not supported on Neuron, so for now we disable it with HF DEACTIVATE ASYNC LOAD=1 . The command below launches training with torchrun , which starts one process per NeuronCore. TP SIZE and FSDP SIZE set the parallelism layout, and their product is the number of processes: here 4 x 8 = 32, matching --nproc per node . The 32 processes form a 4 x 8 grid. Tensor parallelism splits each layer across the 4 cores of a TP group, and FSDP shards the parameters, gradients and optimizer state across the 8 TP groups. TP SIZE=4 FSDP SIZE=8 HF DEACTIVATE ASYNC LOAD=1 uv run torchrun --nproc per node=32 sft neuron example.py \ --model name or path Qwen/Qwen3-8B \ --dataset name HuggingFaceH4/ultrachat 200k \ --dataset train split train sft \ --output dir Qwen3-8B-sft-tp4-fsdp8 \ --dtype bfloat16 \ --max length 2048 \ --packing true \ --packing strategy wrapped \ --learning rate 2e-5 \ --warmup steps 0.03 \ --num train epochs 1 \ --logging steps 2 \ --per device train batch size 4 \ --gradient accumulation steps 4 \ --save only model true \ --torch compile true We ran it on 8 Trainium2 chips 32 NeuronCores , which is half a trn2.48xlarge instance. Training runs to completion, and trainer.save model writes a standard Transformers checkpoint. Nothing in it is tied to Neuron: you can push it to the Hugging Face Hub as is and load it on any other hardware. Diffusers Inference Many pipelines from the Diffusers library are compatible with TorchNeuron in eager and compiled mode. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3 are also compatible with tensor parallelism to make the most of your Neuron devices. The following is an example showing how you can use TP: python import torch, torch.distributed as dist, torch neuronx TP from torch.distributed.device mesh import DeviceMesh TP from diffusers import QwenImage21Pipeline, QwenImage21Transformer2DModel, TensorParallelConfig Distributed devices initialization dist.init process group backend="neuron" TP dev = torch.neuron.current device mesh = DeviceMesh "neuron", list range 8 TP=8 Pipeline initialization model = "Qwen/Qwen-Image-2.1" transformer = QwenImage21Transformer2DModel.from pretrained model, subfolder="transformer", torch dtype=torch.bfloat16, parallel config=TensorParallelConfig mesh=mesh , TP pipe = QwenImage21Pipeline.from pretrained model, transformer=transformer, torch dtype=torch.bfloat16, pipe.vae.to dev pipe.transformer.compile repeated blocks backend="neuron", fullgraph=True, dynamic=False Image generation emb, mask, = pipe.encode prompt prompt="A capybara wearing a wizard hat, oil painting", device="cpu" text encoder stays on CPU image = pipe prompt embeds=emb.to dev , prompt embeds mask=mask.to dev if mask is not None else None, height=1024, width=1024, num inference steps=40, generator=torch.Generator "cpu" .manual seed 0 , same noise on every rank .images 0 if dist.get rank == 0: image.save "out.png" You can run this piece of code with: NEURON LOGICAL NC CONFIG=2 torchrun --nproc per node=8 run.py . More and more pipleines will be compatible with TP. If a pipeline you like is not, you can also add support for it following this guide https://huggingface.co/docs/diffusers/main/en/training/distributed inference writing-a-tensor-parallelism-plan . Kernels Make sure to check out the kernels from the Neuron team https://huggingface.co/aws-neuron/kernels as more and more Neuron kernels will become available on the Hub to speed up your models. Neuron kernels can be written using the Neuron Kernel Interface https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/index.html NKI , which gives lower control to optimize operations. This way, users can run their own custom operations and share them on the hub. You can use for example this kernel https://huggingface.co/kernels/aws-neuron/vit-encoder-neuron-kernels to speed up your vision transformers with up to a x2.4 speedup on Trainium Acknowledgments Special thanks to the AWS team for their work on TorchNeuron and for their reactivity to our feedback to make the HF integration as good as possible.