Diffusers documentation
Compilar e descarregar modelos quantizados
Compilar e descarregar modelos quantizados
A otimização de modelos geralmente envolve compensações entre velocidade de inferência e uso de memória. Por exemplo, embora o caching possa aumentar a velocidade de inferência, ele também aumenta o consumo de memória, pois precisa armazenar as saídas das camadas intermediárias de atenção. Uma estratégia de otimização mais equilibrada combina a quantização de um modelo, torch.compile e diversos métodos de offloading.
Consulte o guia de torch.compile para saber mais sobre compilação e como ela pode ser aplicada aqui. Por exemplo, a compilação regional pode reduzir significativamente o tempo de compilação sem abrir mão de nenhuma aceleração.
Para geração de imagens, a combinação de quantização e model offloading pode frequentemente oferecer o melhor equilíbrio entre qualidade, velocidade e memória. O group offloading não é tão eficaz para geração de imagens, pois geralmente não é possível sobrepor totalmente a transferência de dados se o kernel de computação terminar mais rápido. Isso resulta em algum overhead de comunicação entre a CPU e a GPU.
Para geração de vídeo, combinar quantização e group-offloading tende a ser melhor, pois os modelos de vídeo são mais limitados por computação.
A tabela abaixo fornece uma comparação das combinações de estratégias de otimização e seu impacto na latência e no uso de memória para o Flux.
| combinação | latência (s) | uso de memória (GB) |
|---|---|---|
| quantização | 32.602 | 14.9453 |
| quantização, torch.compile | 25.847 | 14.9448 |
| quantização, torch.compile, offloading do modelo para CPU | 32.312 | 12.2369 |
Este guia mostrará como compilar e descarregar um modelo quantizado com bitsandbytes. Certifique-se de estar usando PyTorch nightly e a versão mais recente do bitsandbytes.
Embora usemos bitsandbytes neste exemplo, outros backends de quantização, como TorchAO, também suportam esses recursos.
pip install -U bitsandbytes
Quantização e torch.compile
Comece quantizando um modelo para reduzir a memória necessária para armazenamento e compilando o modelo para acelerar a inferência.
Configure o Dynamo com capture_dynamic_output_shape_ops = True para lidar com saídas dinâmicas ao compilar modelos bitsandbytes.
import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig
torch._dynamo.config.capture_dynamic_output_shape_ops = True
# quantize
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
components_to_quantize=["transformer", "text_encoder_2"],
)
pipeline = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
quantization_config=pipeline_quant_config,
dtype=torch.bfloat16,
).to("cuda") # or "mps", "xpu", "cpu"
# compile
pipeline.transformer.to(memory_format=torch.channels_last)
pipeline.transformer.compile(mode="max-autotune", fullgraph=True)
pipeline("""
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
"""
).images[0]Quantização, torch.compile e offloading
Além da quantização e do torch.compile, experimente o offloading se precisar reduzir ainda mais o uso de memória. O offloading move várias camadas ou componentes do modelo da CPU para a GPU conforme necessário para os cálculos.
Configure o Dynamo cache_size_limit durante o offloading para evitar recompilações excessivas e defina capture_dynamic_output_shape_ops = True para lidar com saídas dinâmicas ao compilar modelos bitsandbytes.
Model CPU offloading move um componente individual do pipeline, como o modelo transformer, para a GPU quando ele é necessário para o cálculo. Caso contrário, ele é transferido para a CPU.
import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig
torch._dynamo.config.cache_size_limit = 1000
torch._dynamo.config.capture_dynamic_output_shape_ops = True
# quantize
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
components_to_quantize=["transformer", "text_encoder_2"],
)
pipeline = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
quantization_config=pipeline_quant_config,
dtype=torch.bfloat16,
).to("cuda") # or "mps", "xpu", "cpu"
# model CPU offloading
pipeline.enable_model_cpu_offload()
# compile
pipeline.transformer.compile()
pipeline(
"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"
).images[0]Group offloading move as camadas internas de um componente individual do pipeline, como o modelo transformer, para a GPU para computação e as descarrega quando não são necessárias. Ao mesmo tempo, utiliza o recurso CUDA stream para pré-carregar a próxima camada para execução.
Ao sobrepor o cálculo e a transferência de dados, é mais rápido que o offloading do modelo para a CPU, ao mesmo tempo em que economiza memória.
# pip install ftfy
import torch
from diffusers import AutoModel, DiffusionPipeline
from diffusers.hooks import apply_group_offloading
from diffusers.utils import export_to_video
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import UMT5EncoderModel
torch._dynamo.config.cache_size_limit = 1000
torch._dynamo.config.capture_dynamic_output_shape_ops = True
# quantize
pipeline_quant_config = PipelineQuantizationConfig(
quant_backend="bitsandbytes_4bit",
quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
components_to_quantize=["transformer", "text_encoder"],
)
text_encoder = UMT5EncoderModel.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers", subfolder="text_encoder", dtype=torch.bfloat16
)
pipeline = DiffusionPipeline.from_pretrained(
"Wan-AI/Wan2.1-T2V-14B-Diffusers",
quantization_config=pipeline_quant_config,
dtype=torch.bfloat16,
).to("cuda") # or "mps", "xpu", "cpu"
# group offloading
onload_device = torch.device("cuda")
offload_device = torch.device("cpu")
pipeline.transformer.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True,
non_blocking=True
)
pipeline.vae.enable_group_offload(
onload_device=onload_device,
offload_device=offload_device,
offload_type="leaf_level",
use_stream=True,
non_blocking=True
)
apply_group_offloading(
pipeline.text_encoder,
onload_device=onload_device,
offload_type="leaf_level",
use_stream=True,
non_blocking=True
)
# compile
pipeline.transformer.compile()
prompt = """
The camera rushes from far to near in a low-angle shot,
revealing a white ferret on a log. It plays, leaps into the water, and emerges, as the camera zooms in
for a close-up. Water splashes berry bushes nearby, while moss, snow, and leaves blanket the ground.
Birch trees and a light blue sky frame the scene, with ferns in the foreground. Side lighting casts dynamic
shadows and warm highlights. Medium composition, front view, low angle, with depth of field.
"""
negative_prompt = """
Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality,
low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured,
misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards
"""
output = pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
num_frames=81,
guidance_scale=5.0,
).frames[0]
export_to_video(output, "output.mp4", fps=16)