Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /v0.43.0 /ru /optimization /device_map.md

mishig

2 days ago

preview code

download

raw

6.18 kB

# Размещение компонентов модели на нескольких GPU с помощью device_map
Аргумент `device_map` размещает компоненты модели на разных GPU. Используйте его, когда модель не помещается на одном GPU, а ваши GPU не имеют быстрого межпроцессорного соединения. См. [Распределённый вывод](../training/distributed_inference), чтобы сравнить его с другими стратегиями, и [Размещение на устройствах](./memory#device-placement) для дополнительных вариантов `device_map`.
Пример ниже загружает каждый компонент только для той стадии, где он нужен, и использует `device_map`, чтобы распределить крупные компоненты между двумя GPU.

## Кодирование промпта

Установите `device_map="balanced"`, чтобы разместить текстовые энкодеры на отдельных GPU. Используйте `max_memory`, чтобы ограничить объём памяти, используемой каждым GPU. Установите `transformer` и `vae` в `None`, чтобы загрузились только текстовые энкодеры.

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

После получения текстовых эмбеддингов удалите текстовые энкодеры с GPU, чтобы освободить место для трансформера.

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

## Денуазинг с помощью шардированного трансформера

Установите `device_map="auto"`, чтобы Accelerate самостоятельно размещал трансформер. Он заполняет первую GPU, прежде чем переходить к следующей, и использует CPU или диск только в том случае, если на GPU заканчивается память. На двух GPU меньшего размера трансформер распределяется между обеими. Веса, находящиеся на CPU или диске, замедляют инференс.

```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]
> Проверьте `transformer.hf_device_map`, чтобы увидеть, как трансформер распределён по устройствам. Пайплайн, загруженный с `device_map`, также имеет атрибут `hf_device_map`, который показывает, где размещён каждый компонент.
Загрузите конвейер только с трансформером и установите `output_type="latent"`, чтобы он возвращал латенты вместо изображений.

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

## Декодирование латентов

Удалите pipeline и transformer из памяти и загрузите VAE для декодирования латентов. VAE достаточно мал, чтобы поместиться на одном 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")
```

За счёт выборочной загрузки и выгрузки моделей, необходимых на каждом этапе, а также распределения самых крупных моделей по нескольким GPU, вы можете выполнять инференс с большими моделями на потребительских GPU.

Xet Storage Details

Size:: 6.18 kB
Xet hash:: 8a4ea6c0a819c09b693bdc06e65656fef6cb4f62f2a3f19d6451516d2bb0bb1e

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

	# Размещение компонентов модели на нескольких GPU с помощью device_map

	Аргумент `device_map` размещает компоненты модели на разных GPU. Используйте его, когда модель не помещается на одном GPU, а ваши GPU не имеют быстрого межпроцессорного соединения. См. [Распределённый вывод](../training/distributed_inference), чтобы сравнить его с другими стратегиями, и [Размещение на устройствах](./memory#device-placement) для дополнительных вариантов `device_map`.

	Пример ниже загружает каждый компонент только для той стадии, где он нужен, и использует `device_map`, чтобы распределить крупные компоненты между двумя GPU.

	## Кодирование промпта

	Установите `device_map="balanced"`, чтобы разместить текстовые энкодеры на отдельных GPU. Используйте `max_memory`, чтобы ограничить объём памяти, используемой каждым GPU. Установите `transformer` и `vae` в `None`, чтобы загрузились только текстовые энкодеры.

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

	После получения текстовых эмбеддингов удалите текстовые энкодеры с GPU, чтобы освободить место для трансформера.

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

	## Денуазинг с помощью шардированного трансформера

	Установите `device_map="auto"`, чтобы Accelerate самостоятельно размещал трансформер. Он заполняет первую GPU, прежде чем переходить к следующей, и использует CPU или диск только в том случае, если на GPU заканчивается память. На двух GPU меньшего размера трансформер распределяется между обеими. Веса, находящиеся на CPU или диске, замедляют инференс.

	```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]
	> Проверьте `transformer.hf_device_map`, чтобы увидеть, как трансформер распределён по устройствам. Пайплайн, загруженный с `device_map`, также имеет атрибут `hf_device_map`, который показывает, где размещён каждый компонент.

	Загрузите конвейер только с трансформером и установите `output_type="latent"`, чтобы он возвращал латенты вместо изображений.

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

	## Декодирование латентов

	Удалите pipeline и transformer из памяти и загрузите VAE для декодирования латентов. VAE достаточно мал, чтобы поместиться на одном 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")
	```

	За счёт выборочной загрузки и выгрузки моделей, необходимых на каждом этапе, а также распределения самых крупных моделей по нескольким GPU, вы можете выполнять инференс с большими моделями на потребительских GPU.