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

Аргумент device_map размещает компоненты модели на разных GPU. Используйте его, когда модель не помещается на одном GPU, а ваши GPU не имеют быстрого межпроцессорного соединения. См. раздел Распределённый вывод, чтобы сравнить его с другими стратегиями, и раздел Размещение на устройствах для дополнительных вариантов device_map.

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

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

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

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, чтобы освободить место для трансформера.

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 или диске, замедляют инференс.

from diffusers import AutoModel
import torch

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

Проверьте transformer.hf_device_map, чтобы увидеть, как трансформер распределён по устройствам. Пайплайн, загруженный с device_map, также имеет атрибут hf_device_map, который показывает, где размещён каждый компонент.

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

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.

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.

Update on GitHub

Diffusers

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

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

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

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