# AWS Trainium

[TorchNeuron](https://github.com/aws-neuron/torch-neuronx) est le backend PyTorch natif pour Trainium2 et Trainium3. Diffusers s’exécute sur ces instances via le package `torch_neuronx`, qui enregistre un appareil `"neuron"`, de sorte que les pipelines s’exécutent avec le même code que sur n’importe quel autre appareil.

Deux modes d'exécution sont disponibles :

| Mode | Comment activer | Notes |
|---|---|---|
| Eager | `pipe.to("neuron")` | Opérations expédiées une à la fois, asynchrones |
| Compiler | `torch.compile(module, backend="neuron")` | Compilation du graphe entier, formes statiques |

Suivez la [documentation AWS Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/) pour configurer le SDK Neuron et installer `torch_neuronx`. Une fois installé, `import torch` le charge automatiquement et enregistre le périphérique `"neuron"`, de sorte que `pipe.to("neuron")` est la seule modification nécessaire. Ajoutez `import torch_neuronx` uniquement si vous avez désactivé le chargement automatique du backend avec `TORCH_DEVICE_BACKEND_AUTOLOAD=0`.

En mode eager, chaque opération est compilée lors de sa première exécution. En mode compilé, le premier appel compile le graphe. Les deux modes mettent en cache le format exécutable Neuron compilé (NEFF) sous `tmp/neff_cache` par défaut, et ce répertoire est effacé au redémarrage. Définissez `TORCH_NEURONX_NEFF_CACHE_DIR` sur un répertoire persistant pour conserver le cache.

## Mode Eager

FLUX.2-klein-4B tient sur un seul NeuronCore, il faut donc déplacer tout le pipeline sur l’appareil `"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 modèle est trop volumineux pour un seul NeuronCore, ou si vous disposez de plusieurs cœurs et souhaitez une latence plus faible, répartissez les modèles sur les cœurs. Consultez la section [Tensor parallelism](#tensor-parallelism).

## Mode compilé

TorchNeuron enregistre `"neuron"` comme nom de backend pour `torch.compile`, de sorte que les composants sont compilés comme n'importe quelle autre cible `torch.compile`. Le premier appel (phase de préchauffage) est lent car il effectue la compilation ; les appels ultérieurs avec les mêmes formes réutilisent le graphe compilé.

> [!IMPORTANT]
> TorchNeuron exige des *formes statiques*, il faut donc passer `dynamic=False`. Une nouvelle valeur de `height` ou de `width` déclenche une nouvelle compilation pour cette forme. Les formes déjà rencontrées sont réutilisées. Modifier `num_inference_steps` ne provoque pas de re-compilation.

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

Pour les transformers avec des blocs répétés, utilisez `compile_repeated_blocks(...)` pour compiler ces blocs individuellement. Ce bloc est réutilisé et réduit considérablement le temps de compilation.

## Parallélisme de tenseurs

Suivez le guide [Tensor parallelism](./tensor_parallelism) pour le chargement des modèles, les plans de partitionnement et les limitations. Pour adapter son exemple à Neuron, remplacez la configuration du groupe de processus CUDA et des appareils par ce qui suit.

Sur Neuron, initialisez le groupe de processus avec `backend="neuron"` et construisez le maillage avec `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)
```

Passez ce `parallel_config` à la méthode [from_pretrained()](/docs/diffusers/main/fr/api/models/overview#diffusers.ModelMixin.from_pretrained) du transformer et utilisez `device` pour les autres composants du pipeline.

Lancez un processus par NeuronCore. Définissez `--nproc_per_node` sur le nombre de NeuronCores sur lesquels vous souhaitez répartir le modèle.

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