Reducir el uso de memoria

Los modelos de difusión modernos como Flux y Wan tienen miles de millones de parámetros que ocupan mucha memoria en tu hardware para la inferencia. Esto es un desafío porque las GPU comunes a menudo no tienen memoria suficiente. Para superar las limitaciones de memoria, puedes usar más de una GPU (si está disponible), descargar algunos de los componentes del pipeline a la CPU, y más.

Esta guía te mostrará cómo reducir el uso de memoria.

Tenga en cuenta que estas técnicas pueden necesitar ajustes dependiendo del modelo. Por ejemplo, un modelo de difusión basado en un transformer podría no beneficiarse de estas optimizaciones de memoria de la misma manera que un modelo basado en UNet.

Varias GPU

Si tienes acceso a más de una GPU, hay algunas opciones para cargar y distribuir de manera eficiente un modelo grande en tu hardware. Estas funciones son compatibles con la biblioteca Accelerate, así que asegúrate de tenerla instalada primero.

pip install -U accelerate

Puntos de control fragmentados

Cargar puntos de control grandes en varios fragmentos es útil porque los fragmentos se cargan uno a la vez. Esto mantiene el uso de memoria bajo, requiriendo solo suficiente memoria para el tamaño del modelo y el tamaño del fragmento más grande. Recomendamos usar fragmentos cuando el punto de control en fp32 es mayor que 5GB. El tamaño de fragmento predeterminado es 5GB.

Divida un punto de control en save_pretrained() con el parámetro max_shard_size.

from diffusers import AutoModel

unet = AutoModel.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet"
)
unet.save_pretrained("sdxl-unet-sharded", max_shard_size="5GB")

Ahora puede usar el punto de control fragmentado, en lugar del punto de control regular, para ahorrar memoria.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

unet = AutoModel.from_pretrained(
    "username/sdxl-unet-sharded", dtype=torch.float16
)
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    unet=unet,
    dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

Colocación de dispositivos

La colocación de dispositivos es una función experimental y la API puede cambiar. Por el momento, solo se admite la estrategia balanced. Planeamos admitir estrategias de mapeo adicionales en el futuro.

El parámetro device_map controla cómo se distribuyen los componentes del modelo en un pipeline o las capas de un modelo individual entre los dispositivos.

nivel de pipeline

nivel de modelo

Al diseñar tu propio device_map, debe ser un diccionario con el nombre específico de un módulo o capa del modelo y un identificador de dispositivo (un entero para las GPUs, cpu para las CPUs y disk para el disco).

Llama a hf_device_map en un modelo para ver cómo se distribuyen las capas del modelo y luego diseña la tuya.

print(transformer.hf_device_map)
{'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 'cpu', 'single_transformer_blocks.11': 'cpu', 'single_transformer_blocks.12': 'cpu', 'single_transformer_blocks.13': 'cpu', 'single_transformer_blocks.14': 'cpu', 'single_transformer_blocks.15': 'cpu', 'single_transformer_blocks.16': 'cpu', 'single_transformer_blocks.17': 'cpu', 'single_transformer_blocks.18': 'cpu', 'single_transformer_blocks.19': 'cpu', 'single_transformer_blocks.20': 'cpu', 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'}

Por ejemplo, el device_map siguiente coloca single_transformer_blocks.10 hasta single_transformer_blocks.20 en una segunda GPU (1).

import torch
from diffusers import AutoModel

device_map = {
    'pos_embed': 0, 'time_text_embed': 0, 'context_embedder': 0, 'x_embedder': 0, 'transformer_blocks': 0, 'single_transformer_blocks.0': 0, 'single_transformer_blocks.1': 0, 'single_transformer_blocks.2': 0, 'single_transformer_blocks.3': 0, 'single_transformer_blocks.4': 0, 'single_transformer_blocks.5': 0, 'single_transformer_blocks.6': 0, 'single_transformer_blocks.7': 0, 'single_transformer_blocks.8': 0, 'single_transformer_blocks.9': 0, 'single_transformer_blocks.10': 1, 'single_transformer_blocks.11': 1, 'single_transformer_blocks.12': 1, 'single_transformer_blocks.13': 1, 'single_transformer_blocks.14': 1, 'single_transformer_blocks.15': 1, 'single_transformer_blocks.16': 1, 'single_transformer_blocks.17': 1, 'single_transformer_blocks.18': 1, 'single_transformer_blocks.19': 1, 'single_transformer_blocks.20': 1, 'single_transformer_blocks.21': 'cpu', 'single_transformer_blocks.22': 'cpu', 'single_transformer_blocks.23': 'cpu', 'single_transformer_blocks.24': 'cpu', 'single_transformer_blocks.25': 'cpu', 'single_transformer_blocks.26': 'cpu', 'single_transformer_blocks.27': 'cpu', 'single_transformer_blocks.28': 'cpu', 'single_transformer_blocks.29': 'cpu', 'single_transformer_blocks.30': 'cpu', 'single_transformer_blocks.31': 'cpu', 'single_transformer_blocks.32': 'cpu', 'single_transformer_blocks.33': 'cpu', 'single_transformer_blocks.34': 'cpu', 'single_transformer_blocks.35': 'cpu', 'single_transformer_blocks.36': 'cpu', 'single_transformer_blocks.37': 'cpu', 'norm_out': 'cpu', 'proj_out': 'cpu'
}

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

Pase un diccionario que mapee el uso máximo de memoria a cada dispositivo para aplicar un límite. Si un dispositivo no está en max_memory, se ignora y los componentes del pipeline no se distribuirán en él.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

max_memory = {0:"1GB", 1:"1GB"}
pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="balanced",
    max_memory=max_memory
)

Diffusers utiliza la memoria máxima de todos los dispositivos de forma predeterminada, pero si no caben en las GPU, deberá utilizar una sola GPU y descargar en la CPU con los métodos a continuación.

enable_model_cpu_offload() solo funciona en una GPU, pero un modelo muy grande puede no caber en ella
enable_sequential_cpu_offload() puede funcionar, pero es extremadamente lento y además está limitado a una sola GPU

Use el método reset_device_map() para restablecer el device_map. Esto es necesario si desea utilizar métodos como .to(), enable_sequential_cpu_offload() y enable_model_cpu_offload() en un pipeline que estaba mapeado a dispositivos.

pipeline.reset_device_map()

Segmentación de VAE

El slicing de VAE ahorra memoria dividiendo lotes grandes de entradas en un solo lote de datos y procesándolos por separado. Este método funciona mejor al generar más de una imagen a la vez.

Por ejemplo, si estás generando 4 imágenes a la vez, la decodificación aumentaría la memoria de activación pico en 4x. El rebanado de VAE reduce esto decodificando solo 1 imagen a la vez en lugar de las 4 imágenes simultáneamente.

Llama a enable_slicing() para habilitar el VAE segmentado. Puedes esperar un pequeño aumento en el rendimiento al decodificar lotes de varias imágenes y no hay impacto en el rendimiento para lotes de una sola imagen.

import torch
from diffusers import AutoModel, StableDiffusionXLPipeline

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_slicing()
pipeline(["An astronaut riding a horse on Mars"]*32).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

Las clases AutoencoderKLWan y AsymmetricAutoencoderKL no admiten el uso de slices.

Mosaico de VAE

El tiling del VAE ahorra memoria dividiendo una imagen en mosaicos más pequeños y superpuestos, en lugar de procesar la imagen completa de una sola vez. Esto también reduce el uso de memoria pico, ya que la GPU solo procesa un mosaico a la vez.

Llama a enable_tiling() para habilitar el tiling del VAE. La imagen generada puede presentar alguna variación de tono de una tesela a otra porque se decodifican por separado, pero no debería haber costuras evidentes entre las teselas. El tiling se deshabilita para resoluciones inferiores a un límite preestablecido (pero configurable). Por ejemplo, este límite es 512x512 para el VAE en StableDiffusionPipeline.

import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image

pipeline = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.vae.enable_tiling()

init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/img2img-sdxl-init.png")
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
pipeline(prompt, image=init_image, strength=0.5).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

AutoencoderKLWan y AsymmetricAutoencoderKL no admiten el uso de teselas (tiling).

Descarga

Las estrategias de offloading trasladan las capas o modelos que no están activos actualmente a la CPU para evitar aumentar la memoria de la GPU. Estas estrategias pueden combinarse con cuantización y torch.compile para equilibrar la velocidad de inferencia y el uso de memoria.

Consulte la guía Compilar y descargar modelos cuantizados para obtener más detalles.

Descarga en CPU

El offloading de CPU mueve selectivamente los pesos de la GPU a la CPU. Cuando se requiere un componente, se transfiere a la GPU y cuando no se requiere, se mueve a la CPU. Este método funciona en submódulos en lugar de en modelos completos. Ahorra memoria al evitar almacenar el modelo completo en la GPU.

El offloading a CPU reduce drásticamente el uso de memoria, pero también es extremadamente lento porque los submódulos se transfieren varias veces entre dispositivos. A menudo puede ser impráctico debido a su lentitud.

No muevas el pipeline a CUDA antes de llamar a enable_sequential_cpu_offload(), de lo contrario la cantidad de memoria ahorrada será mínima (consulta este issue para más detalles). Esta es una operación con estado que instala hooks en el modelo.

Llama a enable_sequential_cpu_offload() para activarlo en un pipeline.

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_sequential_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

Descarga de modelos

El offloading de modelos mueve modelos completos a la GPU en lugar de mover selectivamente algunas capas o componentes del modelo. Uno de los principales modelos del pipeline, generalmente el codificador de texto, la UNet y el VAE, se coloca en la GPU, mientras que los demás componentes se mantienen en la CPU. Los componentes como la UNet que se ejecutan varias veces permanecen en la GPU hasta que terminan por completo y ya no se necesitan. Esto elimina la sobrecarga de comunicación del offloading de CPU y hace que el offloading de modelos sea una alternativa más rápida. El contrapunto es que el ahorro de memoria no será tan grande.

Tenga en cuenta que si los modelos se reutilizan fuera de la pipeline después de que se hayan instalado los hooks (consulte Removing Hooks para más detalles), debe ejecutar la pipeline completa y los modelos en el orden esperado para poder descargarlos correctamente. Esta es una operación con estado que instala hooks en el modelo.

Llama a enable_model_cpu_offload() para activarlo en un pipeline.

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", dtype=torch.bfloat16
)
pipeline.enable_model_cpu_offload()

pipeline(
    prompt="An astronaut riding a horse on Mars",
    guidance_scale=0.,
    height=768,
    width=1360,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

enable_model_cpu_offload() también resulta útil cuando se utiliza el método encode_prompt() de forma independiente para generar el estado oculto de los codificadores de texto.

Descarga de grupos

El offloading de grupos mueve grupos de capas internas (torch.nn.ModuleList o torch.nn.Sequential) a la CPU. Utiliza menos memoria que el offloading de modelo y es más rápido que el offloading de CPU porque reduce la sobrecarga de comunicación.

El descargado por grupos puede no funcionar con todos los modelos si la implementación de la propagación hacia adelante contiene un cambio de dispositivo de los datos de entrada dependiente de los pesos, ya que puede entrar en conflicto con el mecanismo de cambio de dispositivo del descargado por grupos.

Active el descargado de grupos configurando el parámetro offload_type en block_level o leaf_level.

block_level descarga grupos de capas en función del parámetro num_blocks_per_group. Por ejemplo, si num_blocks_per_group=2 en un modelo con 40 capas, se cargan y descargan 2 capas a la vez (20 cargas/descargas en total). Esto reduce drásticamente los requisitos de memoria.
leaf_level descarga capas individuales en el nivel más bajo y es equivalente a descarga a CPU. Sin embargo, puede hacerse más rápido si se usan streams sin sacrificar la velocidad de inferencia.

El offloading por grupos se admite para pipelines completos o para modelos individuales. Aplicar el offloading por grupos a todo el pipeline es la opción más sencilla, mientras que aplicarlo de forma selectiva a modelos individuales ofrece a los usuarios mayor flexibilidad para utilizar técnicas de offloading diferentes para distintos modelos.

pipeline

model

CUDA stream

El parámetro use_stream puede activarse en dispositivos CUDA que admiten flujos de transferencia de datos asíncronos para reducir el tiempo de ejecución total en comparación con descarga a CPU. Superpone la transferencia de datos y el cálculo mediante la preobtención de capas. La siguiente capa a ejecutar se carga en la GPU mientras la capa actual aún se está ejecutando. Puede aumentar significativamente la memoria de CPU, por lo que asegúrese de contar con el doble de memoria que el tamaño del modelo.

Establezca record_stream=True para obtener una mayor aceleración a costa de un ligero aumento en el uso de memoria. Consulte la documentación de torch.Tensor.record_stream para obtener más información.

Cuando use_stream=True en VAEs con tiling habilitado, asegúrate de realizar una pasada de adelante (forward pass) de prueba (posible también con entradas de prueba) antes de la inferencia para evitar errores de incompatibilidad de dispositivo. Esto puede no funcionar en todas las implementaciones, así que no dudes en abrir un problema (issue) si encuentras algún inconveniente.

Si está utilizando el descargado de grupos a nivel de bloque block_level con use_stream habilitado, el parámetro num_blocks_per_group debe establecerse en 1; de lo contrario, se emitirá una advertencia.

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", use_stream=True, record_stream=True)

El parámetro low_cpu_mem_usage se puede establecer en True para reducir el uso de memoria de CPU al usar streams durante el offloading de grupos. Es lo más adecuado para el offloading de leaf_level y cuando la memoria de CPU es el cuello de botella. Se ahorra memoria creando tensores fijos sobre la marcha en lugar de fijarlos previamente. Sin embargo, esto puede aumentar el tiempo de ejecución total.

Descarga a disco

El descargado de grupos puede consumir una cantidad significativa de memoria del sistema dependiendo del tamaño del modelo. En sistemas con memoria limitada, intente realizar el descargado de grupos en el disco como memoria secundaria.

Establezca el argumento offload_to_disk_path en enable_group_offload() o apply_group_offloading() para descargar el modelo en el disco.

pipeline.transformer.enable_group_offload(onload_device=onload_device, offload_device=offload_device, offload_type="leaf_level", offload_to_disk_path="path/to/disk")

apply_group_offloading(pipeline.text_encoder, onload_device=onload_device, offload_type="block_level", num_blocks_per_group=2, offload_to_disk_path="path/to/disk")

Consulte estas dos tablas para comparar los compromisos entre velocidad y memoria.

Casting por capa

Combine el casteo por capa con group offloading para ahorrar aún más memoria.

El casteo por capas almacena los pesos en un formato de datos más pequeño (por ejemplo, torch.float8_e4m3fn y torch.float8_e5m2) para usar menos memoria y los convierte a una precisión más alta, como torch.float16 o torch.bfloat16, para el cálculo. Se omiten ciertas capas (los pesos relacionados con la normalización y la modulación) porque almacenarlos en fp8 puede degradar la calidad de la generación.

El casteo por capas puede no funcionar con todos los modelos si la implementación de la función forward contiene un casteo interno de pesos. La implementación actual del casteo por capas asume que el paso forward es independiente de la precisión de los pesos y que los tipos de datos de entrada siempre se especifican en compute_dtype (ver aquí para una implementación incompatible).

El casteo por capas también puede fallar en implementaciones de modelado personalizadas con capas de PEFT. Existen algunas comprobaciones disponibles, pero no se han probado exhaustivamente ni se garantiza que funcionen en todos los casos.

Llama a enable_layerwise_casting() para establecer los tipos de datos de almacenamiento y de cómputo.

import torch
from diffusers import CogVideoXPipeline, CogVideoXTransformer3DModel
from diffusers.utils import export_to_video

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)
transformer.enable_layerwise_casting(storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)

pipeline = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b",
    transformer=transformer,
    dtype=torch.bfloat16
).to("cuda")  # or "mps", "xpu", "cpu"
prompt = (
    "A panda, dressed in a small, red jacket and a tiny hat, sits on a wooden stool in a serene bamboo forest. "
    "The panda's fluffy paws strum a miniature acoustic guitar, producing soft, melodic tunes. Nearby, a few other "
    "pandas gather, watching curiously and some clapping in rhythm. Sunlight filters through the tall bamboo, "
    "casting a gentle glow on the scene. The panda's face is expressive, showing concentration and joy as it plays. "
    "The background includes a small, flowing stream and vibrant green foliage, enhancing the peaceful and magical "
    "atmosphere of this unique musical performance."
)
video = pipeline(prompt=prompt, guidance_scale=6, num_inference_steps=50).frames[0]
print(f"Max memory reserved: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")
export_to_video(video, "output.mp4", fps=8)

El método apply_layerwise_casting() también se puede utilizar si se necesita un mayor control y flexibilidad. Puede aplicarse parcialmente a las capas del modelo llamándolo en módulos internos específicos. Utilice los parámetros skip_modules_pattern o skip_modules_classes para especificar los módulos que se deben evitar, como las capas de normalización y modulación.

import torch
from diffusers import CogVideoXTransformer3DModel
from diffusers.hooks import apply_layerwise_casting

transformer = CogVideoXTransformer3DModel.from_pretrained(
    "THUDM/CogVideoX-5b",
    subfolder="transformer",
    dtype=torch.bfloat16
)

# skip the normalization layer
apply_layerwise_casting(
    transformer,
    storage_dtype=torch.float8_e4m3fn,
    compute_dtype=torch.bfloat16,
    skip_modules_classes=["norm"],
    non_blocking=True,
)

torch.channels_last

torch.channels_last invierte cómo se almacenan los tensores, pasando de (batch size, channels, height, width) a (batch size, height, width, channels). Esto alinea los tensores con la forma en que el hardware accede secuencialmente a los tensores almacenados en memoria y evita saltar por la memoria para acceder a los valores de píxel.

No todos los operadores admiten actualmente el formato channels-last y pueden dar lugar a un rendimiento inferior, pero aun así vale la pena intentarlo.

print(pipeline.unet.conv_out.state_dict()["weight"].stride())  # (2880, 9, 3, 1)
pipeline.unet.to(memory_format=torch.channels_last)  # in-place operation
print(
    pipeline.unet.conv_out.state_dict()["weight"].stride()
)  # (2880, 1, 960, 320) having a stride of 1 for the 2nd dimension proves that it works

Atención eficiente en memoria

Diffusers admite varios backends de atención eficientes en memoria (FlashAttention, xFormers, SageAttention, entre otros) a través de set_attention_backend(). Consulta la guía de Attention backends para aprender a cambiar entre ellos.

Update on GitHub