AWS Trainium

TorchNeuron 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 :

Suivez la documentation AWS Neuron 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".

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.

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

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.

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 pour le chargement des modèles, les plans de sharding et les limitations. Pour adapter son exemple à Neuron, remplacez la configuration du process-group CUDA et des appareils par la suivante.

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

Passez ce parallel_config à la méthode 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.

torchrun --nproc_per_node=8 tensor_parallel_neuron.py