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

What it claims

Hugging Face now recommends that Trainium and Inferentia users run Transformers, Diffusers, Accelerate, PEFT and TRL directly on TorchNeuron instead of going through Optimum Neuron. Optimum Neuron sat on torch-neuronx's PyTorch XLA path, which meant lazy tensors, static shapes only and a recompile on every graph break. It kept custom modeling classes and trainers to hide that. TorchNeuron registers neuron as a native PyTorch device, so the standard HF classes and scripts run unchanged. The public beta starts 20 October 2026.

Evidence

The post reports that 89% of the 100 most-used text-generation architectures work out of the box with TorchNeuron, against 46% with Optimum Neuron v0.4.5. The SFT example ran to completion on 8 Trainium2 chips (32 NeuronCores, half a trn2.48xlarge). It produced a standard Transformers checkpoint that loads on any hardware. The post gives no throughput, cost or loss figures for any of these runs. A linked NKI kernel claims up to 2.4x speedup for vision transformers on Trainium, with no baseline or setting stated.

How to use it

pip install torch-neuronx

For inference, pass device_map="neuron". For tensor parallelism, initialise dist.init_process_group(backend="neuron") and pass DistributedConfig(tp_size=world_size) plus a DeviceMesh("neuron", ...) to from_pretrained. The post strongly recommends cache_implementation="static" with a fixed max_cache_len and prefill_chunk_size, inputs padded to a multiple of the chunk size, and CompileConfig(backend="neuron", fullgraph=False, dynamic=False). All of these avoid recompilation.

For training, a plain TRL SFTTrainer script works. Parallelism goes in model_init_kwargs as DistributedConfig(tp_size=..., fsdp_size=...). The model is sharded as weights load, and the Trainer does not re-wrap it, so no --fsdp flag or Accelerate FSDP config is needed. The Neuron-specific settings are:

• torch._dynamo.config.automatic_dynamic_shapes = False
• packing with --packing_strategy wrapped and a fixed --max_length
• --torch_compile true, which selects the Neuron backend automatically
• HF_DEACTIVATE_ASYNC_LOAD=1

With TRL 1.16 or later, Neuron batches are padded to multiples of 128 tokens. The run command is a 4 x 8 grid, where TP splits each layer across 4 cores and FSDP shards 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

In Diffusers, FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1 and H3 support TP through TensorParallelConfig(mesh=mesh). The example compiles with pipe.transformer.compile_repeated_blocks(backend="neuron", fullgraph=True, dynamic=False) and keeps the text encoder on CPU. It launches with:

NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py

Custom kernels can be written in the Neuron Kernel Interface (NKI) and shared on the Hub.

Limits

TorchNeuron is a public beta. Dynamic shapes are still a problem. The Neuron backend cannot compile Dynamo's automatic dynamic shapes (TRL issue #7660), so every example pins shapes to keep compiled graphs reusable. Asynchronous weight loading in Transformers is unsupported. Only the five named Diffusers pipelines support TP. Other pipelines need porting by following a linked guide. The 11% of top text-generation architectures that do not work are not named.