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

TorchNeuron provides a new native backend for PyTorch on AWS Trainium. Transformers, Diffusers and Kernels natively support it, making it easier than ever to run models on Trainium and to share them on the HF hub.

TorchNeuron

TorchNeuron 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 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 . 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.

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.