Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /main /fr /optimization /device_map.md

mishig

3 days ago

preview code

download

raw

4.84 kB

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

Xet Storage Details

Size:: 4.84 kB
Xet hash:: 6aa78cb84f6ef44a46010f66cfc8ecc0a7cc96163238cbbad34ce1597ada4dd8

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

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