Distribuir los componentes del modelo entre GPUs con device_map

El argumento device_map coloca los componentes del modelo en GPUs separadas. Úsalo cuando un modelo no quepa en una sola GPU y tus GPUs no cuenten con un interconexión rápida. Consulta Inferencia distribuida para compararlo con las demás estrategias, y Colocación de dispositivos para más opciones de device_map.

El ejemplo siguiente carga cada componente solo para la etapa que lo necesita y usa device_map para distribuir los más grandes entre dos GPUs.

Codificar el prompt

Establezca device_map="balanced" para colocar los codificadores de texto en GPUs separadas. Use max_memory para limitar la cantidad de memoria que utiliza cada GPU. Establezca transformer y vae en None para que solo se carguen los codificadores de texto.

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
    )

Una vez que tengas los embeddings de texto, elimina los codificadores de texto de la GPU para hacer espacio para el 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()

Denoise con el transformer fragmentado

Establezca device_map="auto" para que Accelerate coloque el transformer. Llena la primera GPU antes de pasar a la siguiente y solo utiliza la CPU o el disco cuando las GPUs se quedan sin memoria. En dos GPUs más pequeñas, el transformer queda dividido entre ambas. Los pesos en la CPU o en el disco ralentizan la inferencia.

from diffusers import AutoModel
import torch

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

Compruebe transformer.hf_device_map para ver cómo se distribuye el transformer entre los dispositivos. Un pipeline cargado con device_map también tiene un atributo hf_device_map que muestra dónde se coloca cada componente.

Carga un pipeline con solo el transformer y establece output_type="latent" para que devuelva latentes en lugar de imágenes.

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

Decodificar los latentes

Elimine el pipeline y el transformer de la memoria y cargue un VAE para decodificar los latentes. El VAE es lo suficientemente pequeño como para caber en una sola 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")

Al cargar y descargar selectivamente los modelos que necesita en cada etapa y dividir los modelos más grandes entre varias GPU, puede ejecutar inferencia con modelos grandes en GPU de consumo.

Update on GitHub

Diffusers

Distribuir los componentes del modelo entre GPUs con device_map

Codificar el prompt

Denoise con el transformer fragmentado

Decodificar los latentes