# AWS Trainium

[TorchNeuron](https://github.com/aws-neuron/torch-neuronx) 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:

| Modo | Cómo activar | Notas |
|---|---|---|
| Eager | `pipe.to("neuron")` | Operaciones despachadas una a la vez, asíncronas |
| Compilar | `torch.compile(module, backend="neuron")` | Compilación de todo el grafo, formas estáticas |

Siga la [documentación de AWS Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/) 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"`.

```python
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 deseas una latencia más baja, divide los modelos entre los núcleos. Consulta la sección [Tensor parallelism](#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.

> [!IMPORTANT]
> 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.

```python
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](../training/distributed_inference#tensor-parallelism) 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", ...)`.

```python
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()](/docs/diffusers/main/es/api/models/overview#diffusers.ModelMixin.from_pretrained) del transformer y usa `device` para los componentes restantes del pipeline.

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

```bash
torchrun --nproc_per_node=8 tensor_parallel_neuron.py
```