TorchNeuron: operating HF fashions on AWS Trainium is less complicated than ever

TorchNeuron offers a brand new native backend for PyTorch on AWS Trainium. Transformers, Diffusers and Kernels natively help it, making it simpler than ever to run fashions on Trainium and to share them on the HF hub.

TorchNeuron

TorchNeuron is a PyTorch plugin that brings a brand new backend for Trainium enabling help for keen mode, torch.compile and native distributed APIs like FSDP and TP.
Learn the Pytorch weblog to study extra about TorchNeuron.

Beginning Oct. 20, 2026, TorchNeuron can be accessible as a public beta. Researchers and builders will have the ability to set up it with pip set up torch-neuronx.

New backend, native interface in HF libraries

Thus far, the Hugging Face expertise on AWS Trainium and Inferentia — AI chips designed respectively for cost-effective coaching and inference at scale — has been delivered by way of optimum-neuron. This library depends on the torch-neuronx bundle, which relies on PyTorch XLA. It was the pure alternative on the time, however it meant considering like a compiler: lazy tensors as a substitute of keen execution, static shapes solely, a recompilation for each graph break. Optimum Neuron existed to cover that complexity, and maintained customized code (modeling courses, trainers, and many others) to take action.

TorchNeuron modifications this. It integrates straight into PyTorch and exposes neuron as a local machine, with keen execution, PyTorch’s personal distributed primitives (DTensor, FSDP2) and torch.compile.

For the Hugging Face libraries, Trainium turns into a tool like some other: Transformers, Diffusers, Speed up, PEFT and TRL run on it straight, with the identical courses and scripts you’ll use on some other accelerator. Any longer, we advocate that Neuron customers use these libraries straight as a substitute of going by way of Optimum Neuron.

Transformers

Inference

Textual content era is now natively supported in Transformers for a variety of fashions. With TorchNeuron, 89% of the highest 100 most-used text-generation mannequin architectures are supported out of the field, in comparison with 46% with Optimum Neuron v0.4.5. You simply have to move the device_map=neuron argument in your code and your mannequin can be moved to the Neuron units.

The next is an instance of text-generation with tensor parallelism, which can allow to completely profit out of your Neuron occasion. Observe that it makes use of a static cache for era as that is strongly advisable with Neuron to stop frequent recompilations.

import os

import torch
import torch.distributed as dist
import torch_neuronx  
from torch.distributed import DeviceMesh
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
from transformers.distributed import DistributedConfig
from transformers.era.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()

mannequin = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    distributed_config=DistributedConfig(tp_size=world_size),
    device_mesh=DeviceMesh("neuron", checklist(vary(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 = mannequin.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        generation_config=generation_config,
    )

if dist.get_rank() == 0:
    prompt_len = inputs["input_ids"].form[1]
    print(tokenizer.decode(output_ids[0, prompt_len:], skip_special_tokens=True))

Coaching

Now that Neuron is a local PyTorch machine, coaching on Trainium makes use of the identical Hugging Face libraries as some other accelerator: Transformers, Speed up, PEFT and TRL. You not want a devoted coach or a Neuron-specific fork of the mannequin code. A regular TRL fine-tuning script runs on a Trainium occasion with only some Neuron-specific settings.

Here’s a supervised fine-tuning script constructed on TRL’s SFTTrainer. If in case you have fine-tuned a mannequin with TRL on GPUs earlier than, it ought to look acquainted.

import os

import torch
from datasets import load_dataset
from transformers.distributed import DistributedConfig

from trl import ModelConfig, ScriptArguments, SFTConfig, SFTTrainer, TrlParser


def primary(script_args, training_args, model_args):
    
    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:
        
        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,
        title=script_args.dataset_config,
        break up=script_args.dataset_train_split,
    )
    
    train_dataset = train_dataset.select_columns(["messages"])

    
    coach = SFTTrainer(
        mannequin=model_args.model_name_or_path,
        args=training_args,
        train_dataset=train_dataset,
    )

    coach.practice()
    coach.save_model(training_args.output_dir)
    coach.finish()


if __name__ == "__main__":
    parser = TrlParser((ScriptArguments, SFTConfig, ModelConfig))
    script_args, training_args, model_args = parser.parse_args_and_config()
    primary(script_args, training_args, model_args)

Parallelism is dealt with natively by Transformers. You place a DistributedConfig with the tensor-parallel (TP) and FSDP sizes in model_init_kwargs, SFTTrainer passes it to from_pretrained, and the mannequin is sharded as its weights are loaded. The Coach sees that the mannequin is already distributed and doesn’t wrap it once more. You do not want an --fsdp flag or an Speed up FSDP config.

Nothing within the script is tied to Trainium: the identical code runs on any {hardware}. A couple of issues are value realizing if you run it on Trainium:

• Static shapes. The script disables automated dynamic shapes in torch.compile and makes use of wrapped packing to supply sequences of as much as 2048 tokens. With TRL ≥ 1.16, batches on Neuron are padded to a a number of of 128 tokens, limiting the variety of enter shapes. Compiled graphs are reused for subsequent steps with the identical shapes.
• torch.compile on. Passing --torch_compile true is all it takes: on a Trainium occasion, torch.compile detects the machine and picks the Neuron backend mechanically, with nothing else to configure.
• Async weight loading off. The built-in asynchronous weight loading in Transformers shouldn’t be supported on Neuron, so for now we disable it with HF_DEACTIVATE_ASYNC_LOAD=1.

The command beneath launches coaching with torchrun, which begins one course of per NeuronCore. TP_SIZE and FSDP_SIZE set the parallelism structure, and their product is the variety of processes: right here 4 x 8 = 32, matching --nproc_per_node. The 32 processes type a 4 x 8 grid. Tensor parallelism splits every layer throughout the 4 cores of a TP group, and FSDP shards the parameters, gradients and optimizer state throughout the 8 TP teams.

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

Coaching runs to completion, and coach.save_model writes a normal Transformers checkpoint. Nothing in it’s tied to Neuron: you’ll be able to push it to the Hugging Face Hub as is and cargo it on some other {hardware}.

Diffusers

Inference

Many pipelines from the Diffusers library are suitable with TorchNeuron in keen and compiled mode. FLUX.1-dev, FLUX.2, Qwen-Picture, Qwen-Picture-2.1 and H3 are additionally suitable with tensor parallelism to take advantage of your Neuron units. The next is an instance displaying how you should use TP:

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


dist.init_process_group(backend="neuron") 
dev = torch.neuron.current_device()
mesh = DeviceMesh("neuron", checklist(vary(8))) 


mannequin = "Qwen/Qwen-Picture-2.1"
transformer = QwenImage21Transformer2DModel.from_pretrained(
    mannequin,
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
    parallel_config=TensorParallelConfig(mesh=mesh), 
)
pipe = QwenImage21Pipeline.from_pretrained(
    mannequin,
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)
pipe.vae.to(dev)
pipe.transformer.compile_repeated_blocks(backend="neuron", fullgraph=True, dynamic=False)


emb, masks, _ = pipe.encode_prompt(immediate="A capybara sporting a wizard hat, oil portray", machine="cpu")  
picture = pipe(
    prompt_embeds=emb.to(dev),
    prompt_embeds_mask=masks.to(dev) if masks is not None else None,
    peak=1024,
    width=1024,
    num_inference_steps=40,
    generator=torch.Generator("cpu").manual_seed(0), 
).photographs[0]
if dist.get_rank() == 0:
    picture.save("out.png")

You may run this piece of code with: NEURON_LOGICAL_NC_CONFIG=2 torchrun --nproc_per_node=8 run.py.

Increasingly more pipleines can be suitable with TP. If a pipeline you want shouldn’t be, you may also add help for it following this information.

Kernels

Be certain to take a look at the kernels from the Neuron group as an increasing number of Neuron kernels will turn into accessible on the Hub to hurry up your fashions.
Neuron kernels could be written utilizing the Neuron Kernel Interface (NKI), which supplies decrease management to optimize operations. This fashion, customers can run their very own customized operations and share them on the hub.

You need to use for instance this kernel to hurry up your imaginative and prescient transformers with as much as a x2.4 speedup on Trainium!

Acknowledgments

Particular due to the AWS group for his or her work on TorchNeuron and for his or her reactivity to our suggestions to make the HF integration nearly as good as doable.