# Répartir les composants du modèle sur plusieurs GPU avec device_map

L'argument `device_map` place les composants du modèle sur des GPU distincts. Utilisez-le lorsque le modèle ne tient pas sur un seul GPU et que vos GPU ne disposent pas d'un interconnect rapide. Reportez-vous à [Inférence distribuée](../training/distributed_inference) pour le comparer aux autres stratégies, et à [Placement des appareils](./memory#device-placement) pour plus d'options de `device_map`.

L’exemple ci-dessous charge chaque composant uniquement pour l’étape qui en a besoin et utilise `device_map` pour répartir les plus volumineux sur deux GPU.

## Encoder le prompt

Définissez `device_map="balanced"` pour placer les encodeurs de texte sur des GPU distincts. Utilisez `max_memory` pour limiter la quantité de mémoire utilisée par chaque GPU. Définissez `transformer` et `vae` sur `None` afin que seuls les encodeurs de texte soient chargés.

```py
from diffusers import FluxPipeline
import torch

prompt = """
cinematic film still of a cat sipping a margarita in a pool in Palm Springs, California
highly detailed, high budget hollywood movie, cinemascope, moody, epic, gorgeous, film grain
"""

pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=None,
    vae=None,
    device_map="balanced",
    max_memory={0: "16GB", 1: "16GB"},
    dtype=torch.bfloat16
)
with torch.no_grad():
    print("Encoding prompts.")
    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
        prompt=prompt, prompt_2=None, max_sequence_length=512
    )
```

Une fois que vous avez les embeddings de texte, retirez les encodeurs de texte du GPU pour faire de la place au transformer.

```py
import gc

def flush():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

del pipeline.text_encoder
del pipeline.text_encoder_2
del pipeline.tokenizer
del pipeline.tokenizer_2
del pipeline

flush()
```

## Dénuisancer avec le transformer sharded

Définissez `device_map="auto"` pour laisser Accelerate placer le transformer. Il remplit la première GPU avant de passer à la suivante, et n'utilise le CPU ou le disque que lorsque les GPU sont à court de mémoire. Sur deux GPU plus petits, le transformer se retrouve réparti entre les deux. Les poids sur le CPU ou le disque ralentissent l'inférence.

```py
from diffusers import AutoModel
import torch

transformer = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    device_map="auto",
    dtype=torch.bfloat16
)
```

> [!TIP]
> Vérifiez `transformer.hf_device_map` pour voir comment le transformer est réparti sur les appareils. Un pipeline chargé avec `device_map` possède également un attribut `hf_device_map` qui indique l’emplacement de chaque composant.

Chargez un pipeline avec uniquement le transformer et définissez `output_type="latent"` afin qu'il renvoie des latents au lieu d'images.

```py
pipeline = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    vae=None,
    transformer=transformer,
    dtype=torch.bfloat16
)

print("Running denoising.")
height, width = 768, 1360
latents = pipeline(
    prompt_embeds=prompt_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    num_inference_steps=50,
    guidance_scale=3.5,
    height=height,
    width=width,
    output_type="latent",
).images
```

## Décoder les latents

Supprimez le pipeline et le transformer de la mémoire et chargez un VAE pour décoder les latents. Le VAE est suffisamment petit pour tenir sur une seule GPU.

```py
import torch
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor

del pipeline.transformer
del pipeline
del transformer

flush()

vae = AutoencoderKL.from_pretrained(
    "black-forest-labs/FLUX.1-dev", subfolder="vae", dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"
vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)

with torch.no_grad():
    print("Running decoding.")
    latents = FluxPipeline._unpack_latents(latents, height, width, vae_scale_factor)
    latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor

    image = vae.decode(latents, return_dict=False)[0]
    image = image_processor.postprocess(image, output_type="pil")
    image[0].save("split_transformer.png")
```

En chargeant et déchargeant sélectivement les modèles nécessaires à une étape donnée et en répartissant les plus grands modèles sur plusieurs GPU, vous pouvez exécuter l’inférence avec de grands modèles sur des GPU grand public.