Hugging Face's logo

Diffusers documentation

编译和卸载量化模型

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

编译和卸载量化模型

优化模型通常需要在推理速度和内存占用之间进行权衡。例如，缓存虽然可以提升推理速度，但由于需要存储中间注意力层的输出，也会增加内存消耗。一种更为均衡的优化策略是将模型量化、torch.compile以及多种卸载方法相结合。

请参阅 torch.compile 指南，以了解有关编译的更多信息以及在此处如何应用它们。例如，区域编译可以在不牺牲任何加速效果的情况下显著缩短编译时间。

在图像生成中，将量化与模型卸载相结合，通常能在质量、速度和内存之间取得最佳平衡。组卸载在图像生成中效果不佳，因为如果计算内核完成得更快，通常无法完全重叠数据传输。这会导致 CPU 和 GPU 之间产生一定的通信开销。

对于视频生成，结合量化和 group-offloading 通常效果更好，因为视频模型更受计算限制。

下表提供了 Flux 优化策略组合及其对延迟和内存使用影响的比较。

组合延迟（秒）内存使用量（GB）
量化32.60214.9453
量化，torch.compile25.84714.9448
量化、torch.compile、模型 CPU 卸载32.31212.2369
These results are benchmarked on Flux with a RTX 4090. The transformer and text_encoder components are quantized. Refer to the benchmarking script if you're interested in evaluating your own model.

本指南将演示如何使用 bitsandbytes 编译并卸载量化模型。请确保你使用的是 PyTorch nightly 以及 bitsandbytes 的最新版本。

虽然本示例使用了 bitsandbytes，但其他量化后端（例如 TorchAO）也支持这些功能。

pip install -U bitsandbytes

量化与 torch.compile

首先对模型进行量化以降低存储所需的内存，然后对其进行编译以加速推理。

配置 Dynamo 的 capture_dynamic_output_shape_ops = True，以在编译 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]

量化、torch.compile 和卸载

除了量化和 torch.compile，如果还需要进一步降低内存占用，可以尝试使用 offloading。Offloading 会根据计算需求，将各个层或模型组件在 CPU 和 GPU 之间进行移动。

在卸载过程中配置 Dynamo 的 cache_size_limit 以避免过度重新编译，并设置 capture_dynamic_output_shape_ops = True 以在处理 bitsandbytes 模型的编译时应对动态输出。

模型 CPU 卸载
组卸载

模型 CPU 卸载会在需要计算时，将单个流水线组件（例如 transformer 模型）移至 GPU。否则，该组件会被卸载到 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]
Update on GitHub

←远程推理 概述→