AWS Trainium

TorchNeuron es el backend nativo de PyTorch para Trainium2 y Trainium3. Diffusers se ejecuta en estas instancias a través del paquete torch_neuronx, que registra un dispositivo "neuron", de modo que los pipelines se ejecutan con el mismo código que en cualquier otro dispositivo.

Hay dos modos de ejecución disponibles:

Siga la documentación de AWS Neuron para configurar el SDK de Neuron e instalar torch_neuronx. Una vez instalado, import torch lo carga automáticamente y registra el dispositivo "neuron", por lo que pipe.to("neuron") es el único cambio necesario. Añada import torch_neuronx solo si deshabilitó la carga automática del backend con TORCH_DEVICE_BACKEND_AUTOLOAD=0.

En modo eager, cada operación se compila la primera vez que se ejecuta. En modo compilado, la primera llamada compila el grafo. Ambos modos almacenan en caché el formato ejecutable de Neuron (NEFF) compilado en tmp/neff_cache de forma predeterminada, y ese directorio se borra al reiniciar. Establezca TORCH_NEURONX_NEFF_CACHE_DIR en un directorio persistente para conservar la caché.

Modo eager

FLUX.2-klein-4B cabe en un solo NeuronCore, por lo que mueva todo el pipeline al dispositivo "neuron".

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")

Si un modelo es demasiado grande para un solo NeuronCore, o si tienes varios núcleos y quieres una latencia más baja, divide los modelos entre los núcleos. Consulta la sección Tensor parallelism.

Modo compilado

TorchNeuron registra "neuron" como un nombre de backend para torch.compile, por lo que los componentes se compilan como cualquier otro destino de torch.compile. La primera llamada (calentamiento) es lenta porque compila; las llamadas posteriores con las mismas formas reutilizan el gráfico compilado.

TorchNeuron requiere formas estáticas, por lo que debe pasar dynamic=False. Un nuevo height o width vuelve a compilar para esa forma. Las formas ya vistas se reutilizan. Cambiar num_inference_steps no vuelve a compilar.

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")

Para los transformers con bloques repetidos, use compile_repeated_blocks(...) para compilar esos bloques de forma individual. Este bloque se reutiliza y reduce considerablemente el tiempo de compilación.

Paralelismo de tensores

Siga la guía de Paralelismo de tensores para la carga de modelos, los planes de particionamiento y las limitaciones. Para adaptar su ejemplo a Neuron, reemplace la configuración del grupo de procesos y el dispositivo de CUDA con lo siguiente.

En Neuron, inicializa el grupo de procesos con backend="neuron" y construye la malla con 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)

Pasa este parallel_config al from_pretrained() del transformer y usa device para los componentes restantes de la pipeline.

Inicie un proceso por NeuronCore. Establezca --nproc_per_node en el número de NeuronCores que desea dividir.

torchrun --nproc_per_node=8 tensor_parallel_neuron.py