Reduzir o uso de memória

Modelos de difusão modernos, como Flux e Wan, possuem bilhões de parâmetros que ocupam uma grande quantidade de memória no seu hardware para inferência. Isso é desafiador, pois GPUs comuns frequentemente não possuem memória suficiente. Para superar as limitações de memória, você pode usar mais de uma GPU (se disponível), descarregar alguns dos componentes do pipeline para a CPU, entre outras opções.

Este guia mostrará como reduzir o uso de memória.

Lembre-se de que essas técnicas podem precisar ser ajustadas dependendo do modelo. Por exemplo, um modelo de difusão baseado em transformadores pode não se beneficiar igualmente dessas otimizações de memória em comparação com um modelo baseado em UNet.

Várias GPUs

Se você tiver acesso a mais de uma GPU, há algumas opções para carregar e distribuir um modelo grande de forma eficiente entre seu hardware. Esses recursos são suportados pela biblioteca Accelerate, portanto, certifique-se de que ela esteja instalada primeiro.

pip install -U accelerate

Pontos de verificação fragmentados

Carregar checkpoints grandes em vários shards é útil porque os shards são carregados um de cada vez. Isso mantém o uso de memória baixo, exigindo apenas memória suficiente para o tamanho do modelo e o tamanho do maior shard. Recomendamos o uso de sharding quando o checkpoint em fp32 for maior que 5GB. O tamanho padrão do shard é 5GB.

Divida um checkpoint em save_pretrained() com o 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")

Agora você pode usar o checkpoint fragmentado, em vez do checkpoint regular, para economizar memória.

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"

Colocação de dispositivo

O posicionamento de dispositivos é um recurso experimental e a API pode mudar. No momento, apenas a estratégia balanced é suportada. Planejamos suportar estratégias de mapeamento adicionais no futuro.

O parâmetro device_map controla como os componentes do modelo em um pipeline ou as camadas em um modelo individual são distribuídos entre os dispositivos.

nível de pipeline

nível do modelo

Ao projetar seu próprio device_map, ele deve ser um dicionário com o nome específico de um módulo ou camada do modelo e um identificador de dispositivo (um inteiro para GPUs, cpu para CPUs e disk para disco).

Chame hf_device_map em um modelo para ver como as camadas do modelo são distribuídas e, em seguida, crie o seu próprio.

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 exemplo, o device_map abaixo coloca single_transformer_blocks.10 até single_transformer_blocks.20 em uma 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
)

Passe um dicionário mapeando o uso máximo de memória para cada dispositivo para impor um limite. Se um dispositivo não estiver em max_memory, ele será ignorado e os componentes do pipeline não serão distribuídos para ele.

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
)

O Diffusers usa a memória máxima de todos os dispositivos por padrão, mas, se os modelos não couberem nas GPUs, será necessário usar uma única GPU e transferir para a CPU com os métodos abaixo.

enable_model_cpu_offload() funciona apenas em uma única GPU, mas um modelo muito grande pode não caber nela
enable_sequential_cpu_offload() pode funcionar, mas é extremamente lento e também limitado a uma única GPU

Use o método reset_device_map() para redefinir o device_map. Isso é necessário se você quiser usar métodos como .to(), enable_sequential_cpu_offload() e enable_model_cpu_offload() em um pipeline que foi mapeado para dispositivos.

pipeline.reset_device_map()

VAE slicing

O fatiamento do VAE economiza memória dividindo grandes lotes de entradas em um único lote de dados e processando-os separadamente. Este método funciona melhor ao gerar mais de uma imagem por vez.

Por exemplo, se você estiver gerando 4 imagens de uma vez, a decodificação aumentaria a memória de ativação de pico em 4x. O fatiamento do VAE reduz isso decodificando apenas 1 imagem de cada vez, em vez de todas as 4 imagens de uma vez.

Chame enable_slicing() para habilitar o VAE fatiado. Você pode esperar um pequeno aumento de desempenho ao decodificar lotes de múltiplas imagens e nenhum impacto no desempenho para lotes de imagem única.

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

As classes AutoencoderKLWan e AsymmetricAutoencoderKL não suportam slicing.

Tiling do VAE

O tiling do VAE economiza memória dividindo uma imagem em tiles menores sobrepostos, em vez de processar a imagem inteira de uma vez. Isso também reduz o uso de memória de pico, pois a GPU processa apenas um tile por vez.

Chame enable_tiling() para habilitar o tiling do VAE. A imagem gerada pode apresentar alguma variação de tom de um tile para outro, pois eles são decodificados separadamente, mas não deve haver emendas óbvias entre os tiles. O tiling é desabilitado para resoluções menores que um limite pré-especificado (mas configurável). Por exemplo, esse limite é 512x512 para o VAE em 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 e AsymmetricAutoencoderKL não suportam tiling.

Para ignorar completamente a memória do VAE na sua GPU, codifique e decodifique em um endpoint remoto com inferência remota.

Offloading

As estratégias de offloading movem camadas ou modelos que não estão ativos no momento para a CPU, a fim de evitar o aumento do uso de memória da GPU. Essas estratégias podem ser combinadas com quantização e torch.compile para equilibrar a velocidade de inferência e o uso de memória.

Consulte o guia Compile and offloading quantized models para mais detalhes.

Offloading de CPU

O offloading de CPU move seletivamente os pesos da GPU para a CPU. Quando um componente é necessário, ele é transferido para a GPU e, quando não é necessário, é movido para a CPU. Este método funciona em submódulos, em vez de modelos inteiros. Ele economiza memória ao evitar armazenar o modelo inteiro na GPU.

O offloading para CPU reduz drasticamente o uso de memória, mas também é extremamente lento, pois os submódulos são transferidos várias vezes entre os dispositivos. Muitas vezes, isso pode ser impraticável devido à lentidão.

Não mova o pipeline para a CUDA antes de chamar enable_sequential_cpu_offload(), caso contrário, a quantidade de memória economizada será mínima (consulte esta issue para mais detalhes). Esta é uma operação com estado que instala hooks no modelo.

Chame enable_sequential_cpu_offload() para ativá-lo em um 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")

Deslocamento de modelos

O offloading de modelos move modelos inteiros para a GPU, em vez de mover seletivamente algumas camadas ou componentes do modelo. Um dos principais modelos do pipeline, geralmente o codificador de texto, UNet e VAE, é colocado na GPU, enquanto os outros componentes são mantidos na CPU. Componentes como o UNet, que são executados várias vezes, permanecem na GPU até que estejam completamente finalizados e não sejam mais necessários. Isso elimina o overhead de comunicação do CPU offloading e torna o offloading de modelos uma alternativa mais rápida. A compensação é que as economias de memória não serão tão grandes.

Lembre-se de que, se os modelos forem reutilizados fora do pipeline após a instalação dos hooks (veja Removendo Hooks para mais detalhes), você precisa executar o pipeline inteiro e os modelos na ordem esperada para desativá-los corretamente. Esta é uma operação com estado que instala hooks no modelo.

Chame enable_model_cpu_offload() para ativá-lo em um 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() também ajuda quando você usa o método encode_prompt() isoladamente para gerar o estado oculto dos codificadores de texto.

Descarregamento em grupo

O Group offloading move grupos de camadas internas (torch.nn.ModuleList ou torch.nn.Sequential) para a CPU. Ele usa menos memória que o model offloading e é mais rápido que o CPU offloading porque reduz o overhead de comunicação.

O offloading por grupo pode não funcionar com todos os modelos se a implementação do forward contiver o casting de dispositivos dos inputs dependente dos pesos, pois isso pode conflitar com o mecanismo de casting de dispositivos do offloading por grupo.

Ative o offloading de grupo configurando o parâmetro offload_type para block_level ou leaf_level.

block_level descarrega grupos de camadas com base no parâmetro num_blocks_per_group. Por exemplo, se num_blocks_per_group=2 em um modelo com 40 camadas, 2 camadas são carregadas e descarregadas de cada vez (20 carregamentos/descarregamentos no total). Isso reduz drasticamente os requisitos de memória.
leaf_level descarrega camadas individuais no nível mais baixo e é equivalente ao descarregamento para CPU. No entanto, pode ser acelerado se você usar streams sem abrir mão da velocidade de inferência.

O offloading em grupo é suportado para pipelines inteiros ou para modelos individuais. Aplicar o offloading em grupo a todo o pipeline é a opção mais simples, enquanto aplicá-lo seletivamente a modelos individuais oferece aos usuários mais flexibilidade para usar diferentes técnicas de offloading para diferentes modelos.

pipeline

model

CUDA stream

O parâmetro use_stream pode ser ativado para dispositivos CUDA que suportam streams de transferência de dados assíncronos para reduzir o tempo total de execução em comparação com CPU offloading. Ele sobrepõe a transferência de dados e o cálculo usando prefetching de camadas. A próxima camada a ser executada é carregada na GPU enquanto a camada atual ainda está sendo executada. Isso pode aumentar significativamente a memória da CPU, portanto, certifique-se de ter 2x a quantidade de memória do tamanho do modelo.

Defina record_stream=True para obter mais aceleração, com um custo de uso de memória ligeiramente maior. Consulte a documentação de torch.Tensor.record_stream para saber mais.

Quando use_stream=True em VAEs com tiling habilitado, certifique-se de executar uma passagem de forward fictícia (possível também com entradas fictícias) antes da inferência para evitar erros de incompatibilidade de dispositivo. Isso pode não funcionar em todas as implementações, então sinta-se à vontade para abrir uma issue se encontrar algum problema.

Se você estiver usando offloading de grupo em nível de block_level com use_stream habilitado, o parâmetro num_blocks_per_group deve ser definido como 1; caso contrário, um aviso será emitido.

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

O parâmetro low_cpu_mem_usage pode ser definido como True para reduzir o uso de memória de CPU ao usar streams durante o offloading de grupos. É ideal para offloading de leaf_level e quando a memória de CPU é o gargalo. A memória é economizada criando tensores fixados (pinned) sob demanda, em vez de fixá-los previamente. No entanto, isso pode aumentar o tempo total de execução.

Offloading para o disco

O offloading de grupos pode consumir uma quantidade significativa de memória do sistema, dependendo do tamanho do modelo. Em sistemas com memória limitada, tente realizar o offloading de grupos para o disco como memória secundária.

Defina o argumento offload_to_disk_path em enable_group_offload() ou apply_group_offloading() para descarregar o modelo no 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 duas tabelas para comparar os compromissos entre velocidade e memória.

Casting por camada

Combine o casting por camada com group offloading para economizar ainda mais memória.

O casting por camada armazena os pesos em um formato de dados menor (por exemplo, torch.float8_e4m3fn e torch.float8_e5m2) para usar menos memória e faz o upcast desses pesos para uma precisão maior, como torch.float16 ou torch.bfloat16, para o cálculo. Certas camadas (pesos relacionados a normalização e modulação) são ignoradas, pois armazená-los em fp8 pode degradar a qualidade da geração.

O casting por camada pode não funcionar com todos os modelos se a implementação do forward contiver conversão de tipos interna dos pesos. A implementação atual do casting por camada pressupõe que a passagem forward seja independente da precisão dos pesos e que os tipos de dados de entrada sejam sempre especificados em compute_dtype (veja aqui para uma implementação incompatível).

O casting por camada também pode falhar em implementações de modelagem personalizadas com camadas de PEFT. Existem algumas verificações disponíveis, mas elas não são amplamente testadas nem garantidas para funcionar em todos os casos.

Chame enable_layerwise_casting() para definir os tipos de dados de armazenamento e computação.

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)

O método apply_layerwise_casting() também pode ser usado se você precisar de mais controle e flexibilidade. Ele pode ser aplicado parcialmente às camadas do modelo chamando-o em módulos internos específicos. Use os parâmetros skip_modules_pattern ou skip_modules_classes para especificar os módulos a serem evitados, como as camadas de normalização e modulação.

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 inverte a forma como os tensors são armazenados, de (batch size, channels, height, width) para (batch size, height, width, channels). Isso alinha os tensors com a forma como o hardware acessa sequencialmente os tensors armazenados na memória e evita saltos na memória para acessar os valores dos pixels.

Nem todos os operadores suportam atualmente o formato channels-last e podem resultar em pior desempenho, mas ainda vale a pena tentar.

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

Atenção eficiente em memória

O Diffusers suporta vários backends de atenção eficientes em memória (FlashAttention, xFormers, SageAttention, entre outros) por meio de set_attention_backend(). Consulte o guia Attention backends para aprender a alternar entre eles.

Update on GitHub