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 pour le comparer aux autres stratégies, et à Placement des appareils 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.

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.

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.

from diffusers import AutoModel
import torch

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

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.

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.

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.

Update on GitHub

Diffusers

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

Encoder le prompt

Dénuisancer avec le transformer sharded

Décoder les latents