AWS Trainium

TorchNeuron is the native PyTorch backend for Trainium2 and Trainium3. Diffusers runs on these instances through the torch_neuronx package, which registers a "neuron" device, so pipelines run with the same code as on any other device.

Two execution modes are available:

Follow the AWS Neuron documentation to set up the Neuron SDK and install torch_neuronx. Once installed, import torch loads it automatically and registers the "neuron" device, so pipe.to("neuron") is the only change needed. Add import torch_neuronx only if you disabled backend autoloading with TORCH_DEVICE_BACKEND_AUTOLOAD=0.

In eager mode, each operation compiles the first time it runs. In compiled mode, the first call compiles the graph. Both modes cache the compiled Neuron executable format (NEFF) under tmp/neff_cache by default, and that directory is cleared on reboot. Set TORCH_NEURONX_NEFF_CACHE_DIR to a persistent directory to keep the cache.

Eager mode

FLUX.2-klein-4B fits on a single NeuronCore, so move the whole pipeline to the "neuron" device.

import torch

from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", dtype=torch.bfloat16)
pipe.to("neuron")

image = pipe(
    prompt="a golden retriever surfing a wave, photorealistic",
    height=1024,
    width=1024,
    num_inference_steps=4,
).images[0]

image.save("output.png")

If a model is too large for a single NeuronCore, or you have several cores and want lower latency, shard the models across cores instead. See the Tensor parallelism section.

Compiled mode

TorchNeuron registers "neuron" as a torch.compile backend name, so components compile like any other torch.compile target. The first call (warmup) is slow because it compiles; later calls with the same shapes reuse the compiled graph.

TorchNeuron requires static shapes, so pass dynamic=False. A new height or width compiles again for that shape. Shapes already seen are reused. Changing num_inference_steps doesn’t recompile.

import torch

from diffusers import Flux2KleinPipeline

pipe = Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", dtype=torch.bfloat16)
pipe.to("neuron")
pipe.transformer.compile(backend="neuron", fullgraph=True, dynamic=False)

# Warmup — triggers static graph compilation.
pipe(prompt="warmup", height=1024, width=1024, num_inference_steps=4)

# Later calls with the same shapes reuse the compiled graph.
image = pipe(
    prompt="a golden retriever surfing a wave, photorealistic",
    height=1024,
    width=1024,
    num_inference_steps=4,
).images[0]

image.save("output.png")

For transformers with repeated blocks, use compile_repeated_blocks(...) to compile those blocks individually. This block is reused and reduces compile time considerably.

Tensor parallelism

Follow the Tensor parallelism guide for model loading, sharding plans, and limitations. To adapt its example for Neuron, replace the CUDA process-group and device setup with the following.

On Neuron, initialize the process group with backend="neuron" and build the mesh with DeviceMesh("neuron", ...).

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import DeviceMesh

from diffusers import TensorParallelConfig

dist.init_process_group(backend="neuron")
mesh = DeviceMesh("neuron", list(range(dist.get_world_size())))
device = torch.neuron.current_device()
parallel_config = TensorParallelConfig(mesh=mesh)

Pass this parallel_config to the transformer’s from_pretrained() and use device for the remaining pipeline components.

Launch one process per NeuronCore. Set --nproc_per_node to the number of NeuronCores you want to shard across.

torchrun --nproc_per_node=8 tensor_parallel_neuron.py