Diffusers documentation
Cuantizar, compilar y descargar
Cuantizar, compilar y descargar
Optimizar modelos suele implicar un equilibrio entre la velocidad de inferencia y el uso de memoria. Por ejemplo, aunque la memoria caché puede aumentar la velocidad de inferencia, también incrementa el consumo de memoria, ya que es necesario almacenar las salidas de las capas intermedias de atención. Una estrategia de optimización más equilibrada combina la cuantización de un modelo, torch.compile y diversos métodos de descarga.
Consulte la guía de torch.compile para obtener más información sobre la compilación y cómo puede aplicarse aquí. Por ejemplo, la compilación regional puede reducir significativamente el tiempo de compilación sin renunciar a ninguna aceleración.
Para la generación de imágenes, combinar la cuantización y el descarga del modelo suele ofrecer el mejor equilibrio entre calidad, velocidad y memoria. La descarga por grupos no es tan efectiva para la generación de imágenes porque generalmente no es posible superponer por completo la transferencia de datos si el kernel de cálculo termina antes. Esto genera cierto sobrecoste de comunicación entre la CPU y la GPU.
Para la generación de video, combinar la cuantización y group-offloading suele ser mejor porque los modelos de video están más limitados por el cómputo.
La tabla siguiente ofrece una comparación de las combinaciones de estrategias de optimización y su impacto en la latencia y el uso de memoria para Flux.
| combinación | latencia (s) | uso de memoria (GB) |
|---|---|---|
| cuantización | 32.602 | 14.9453 |
| cuantización, torch.compile | 25.847 | 14.9448 |
| cuantización, torch.compile, descarga del modelo a CPU | 32.312 | 12.2369 |
Esta guía te mostrará cómo compilar y descargar un modelo cuantizado con bitsandbytes. Asegúrate de estar usando PyTorch nightly y la última versión de bitsandbytes.
Aunque en este ejemplo usamos bitsandbytes, otros backends de cuantización, como TorchAO, también admiten estas funciones.
pip install -U bitsandbytes
Cuantización y torch.compile
Comience cuantizando un modelo para reducir la memoria necesaria para su almacenamiento y compilándolo para acelerar la inferencia.
Configure Dynamo con capture_dynamic_output_shape_ops = True para manejar salidas dinámicas al compilar modelos de 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]Cuantización, torch.compile y descarga
Además de la cuantización y torch.compile, prueba el offloading si necesitas reducir aún más el uso de memoria. El offloading transfiere varias capas o componentes del modelo desde la CPU a la GPU según sea necesario para los cálculos.
Configure el Dynamo cache_size_limit durante el offloading para evitar recompilaciones excesivas y establezca capture_dynamic_output_shape_ops = True para manejar salidas dinámicas al compilar modelos de bitsandbytes.
Model CPU offloading mueve un componente individual de la pipeline, como el modelo transformer, a la GPU cuando se necesita para el cálculo. De lo contrario, se descarga a la 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 mueve las capas internas de un componente individual de la pipeline, como el modelo transformer, a la GPU para su cómputo y las descarga cuando no se necesitan. Al mismo tiempo, utiliza la función CUDA stream para preobtener la siguiente capa para su ejecución.
Al superponer el cálculo y la transferencia de datos, es más rápido que el descargado del modelo a la CPU y, al mismo tiempo, ahorra memoria.
# 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)