TorchNeuron simplifies running Hugging Face models on AWS Trainium

TorchNeuron simplifies running Hugging Face models on AWS Trainium

TorchNeuron is a native PyTorch backend for AWS Trainium. Transformers, Diffusers, and Kernels support it natively, allowing models to run on Trainium and remain shareable through the Hugging Face Hub.

TorchNeuron

TorchNeuron is a PyTorch plugin that adds a Trainium backend with support for eager mode, torch.compile, and native distributed APIs such as FSDP and tensor parallelism (TP). TorchNeuron is scheduled to become available as a public beta on October 20, 2026, and can be installed with:

pip install torch-neuronx

A native backend for Hugging Face libraries

Hugging Face support for AWS Trainium and Inferentia has traditionally been provided through Optimum Neuron. That library relies on torch-neuronx, which is based on PyTorch XLA. This approach required users to work with lazy tensors, static shapes, and recompilation after graph breaks. Optimum Neuron concealed much of that complexity through custom modeling classes, trainers, and related code.

TorchNeuron integrates directly with PyTorch and exposes neuron as a native device. It supports eager execution, PyTorch distributed primitives such as DTensor and FSDP2, and torch.compile.

For Hugging Face libraries, Trainium can now be used like other accelerators. Transformers, Diffusers, Accelerate, PEFT, and TRL run with the same classes and scripts used on other hardware. Neuron users can use these libraries directly rather than going through Optimum Neuron.

Transformers

Inference

Transformers now natively supports text generation for a broad range of models. TorchNeuron supports 89% of the 100 most-used text-generation model architectures out of the box, compared with 46% for Optimum Neuron v0.4.5.

Passing device_map=neuron moves the model to Neuron devices. The following example uses tensor parallelism and a static generation cache, which is recommended on Neuron to avoid frequent recompilation.

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

Because Neuron is a native PyTorch device, training on Trainium uses Transformers, Accelerate, PEFT, and TRL in the same way as other accelerators. Dedicated trainers and Neuron-specific model forks are not required. A standard TRL fine-tuning script can run on Trainium with a few device-specific settings.

The following supervised fine-tuning example uses TRL's SFTTrainer:

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 cannot yet compile Dynamo's automatic dynamic shapes.
    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:
        # Shard the model across the TP and FSDP grid while loading weights.
        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")),
        )

    train_dataset = load_dataset(
        script_args.dataset_name,
        name=script_args.dataset_config,
        split=script_args.dataset_train_split,
    )

    # UltraChat also contains a string `prompt` column. Selecting `messages`
    # ensures that TRL treats the data as a conversational dataset.
    train_dataset = train_dataset.select_columns(["messages"])

    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)

Transformers handles parallelism natively. A DistributedConfig containing the tensor-parallel and FSDP sizes is placed in model_init_kwargs. SFTTrainer passes it to from_pretrained, and the model is sharded during weight loading. Because the Trainer receives an already distributed model, it does not wrap the model again. An --fsdp flag and an Accelerate FSDP configuration are not required.

The script is not specific to Trainium and can run on other hardware. The following settings are relevant when using Neuron:

• Static shapes: The script disables automatic dynamic shapes in torch.compile and uses wrapped packing for sequences of up to 2048 tokens. With TRL 1.16 or later, Neuron batches are padded to a multiple of 128 tokens, limiting the number of input shapes. Compiled graphs are reused for later steps with matching shapes.
• torch.compile: Setting --torch_compile true enables compilation. On Trainium, torch.compile detects the device and selects the Neuron backend automatically.
• Asynchronous weight loading: Transformers' built-in asynchronous weight loading is not supported on Neuron, so it is disabled with HF_DEACTIVATE_ASYNC_LOAD=1.

The following command launches training with torchrun, which starts one process per NeuronCore. TP_SIZE and FSDP_SIZE define the parallelism layout. Their product is the process count: 4 x 8 = 32, matching --nproc_per_node. The 32 processes form a 4 x 8 grid. Tensor parallelism splits each layer across the four cores in a TP group, while FSDP shards parameters, gradients, and optimizer state across the eight 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

This configuration was run on eight Trainium2 chips, providing 32 NeuronCores, which is half of a trn2.48xlarge instance.

The training completes and trainer.save_model creates a standard Transformers checkpoint. The checkpoint is not tied to Neuron and can be pushed to the Hugging Face Hub or loaded on other hardware.

Diffusers

Inference

Many Diffusers pipelines are compatible with TorchNeuron in eager and compiled modes. FLUX.1-dev, FLUX.2, Qwen-Image, Qwen-Image-2.1, and H3 also support tensor parallelism, allowing multiple Neuron devices to be used together.

The following example uses tensor parallelism with Qwen-Image-2.1:

import torch
import torch.distributed as dist
import torch_neuronx
from torch.distributed.device_mesh import DeviceMesh
from diffusers import (
    QwenImage21Pipeline,
    QwenImage21Transformer2DModel,
    TensorParallelConfig,
)

## Initialize distributed devices
dist.init_process_group(backend="neuron")
dev = torch.neuron.current_device()
mesh = DeviceMesh("neuron", list(range(8)))  # TP=8

## Initialize the pipeline
model = "Qwen/Qwen-Image-2.1"
transformer = QwenImage21Transformer2DModel.from_pretrained(
    model,
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
    parallel_config=TensorParallelConfig(mesh=mesh),
)
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,
)

## Generate an image
emb, mask, _ = pipe.encode_prompt(
    prompt="A capybara wearing a wizard hat, oil painting",
    device="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),
).images[0]

if dist.get_rank() == 0:
    image.save("out.png")

The script can be launched with:

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

Additional pipelines can support tensor parallelism by following the Diffusers tensor-parallelism guide.

Kernels

The Neuron team provides kernels through the Hugging Face Hub. These kernels can be written with the Neuron Kernel Interface, which provides lower-level control for optimizing operations. Developers can use the interface to create custom operations and share them on the Hub.

For example, the Vision Transformer encoder Neuron kernel reports up to a 2.4x speedup on Trainium.

Acknowledgments

The TorchNeuron work was developed with the AWS team, including feedback intended to improve the Hugging Face integration.