Diffusers documentation
快速入门
快速入门
量化旨在用更少的位数来表示数据,同时尽量保留原始数据的精度。这通常意味着将数据类型转换为用更少的位数表示相同的信息。例如,如果你的模型权重以 32 位浮点数存储,并将其量化为 16 位浮点数,模型大小将减半,从而更易于存储并减少内存占用。较低的精度还可以加速推理,因为使用更少的位数进行计算所需的时间更短。
Diffusers 支持多种量化后端,使 Flux 等大型扩散模型更易用。本指南介绍如何使用 PipelineQuantizationConfig 类,在从预训练或非量化检查点初始化流水线时对其进行量化。
Pipeline 级别的量化
根据您希望对量化配置进行自定义的程度,有两种方式使用 PipelineQuantizationConfig。
- 对于基本使用场景,请定义
quant_backend、quant_kwargs和components_to_quantize参数 - 如需进行细粒度的量化控制,请定义一个
quant_mapping,为各个模型组件提供量化配置。
基础量化
使用以下参数初始化 PipelineQuantizationConfig。
quant_backend指定要使用的量化后端。目前支持的后端包括:bitsandbytes_4bit、bitsandbytes_8bit、gguf、quanto和torchao。quant_kwargs指定要使用的量化参数。
这些
quant_kwargs参数因后端而异。请参阅 量化 API 文档以查看各后端的参数。
components_to_quantize指定要量化的流水线组件。通常,您应该量化计算量最大的组件,例如 transformer。如果流水线包含多个文本编码器(例如 FluxPipeline),则文本编码器是另一个需要考虑量化的组件。下面的示例对 FluxPipeline 中的 T5 文本编码器进行量化,同时保持 CLIP 模型不变。components_to_quantize可以接受一个列表(用于多个模型)或一个字符串(用于单个模型)。
以下示例使用 BitsAndBytesConfig 中的 load_in_4bit、bnb_4bit_quant_type 和 bnb_4bit_compute_dtype 参数加载 bitsandbytes 后端。
import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig
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_quant_config 传递给 from_pretrained() 以量化该 pipeline。
pipe = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
quantization_config=pipeline_quant_config,
dtype=torch.bfloat16,
).to("cuda") # or "mps", "xpu", "cpu"
image = pipe("photo of a cute dog").images[0]高级量化
quant_mapping 参数提供了更多选项,用于指定流水线中各个组件的量化方式,例如组合使用不同的量化后端。
初始化 PipelineQuantizationConfig 并向其传入 quant_mapping。quant_mapping 允许你为流水线中的每个组件(例如 transformer 和 text encoder)指定量化选项。
下面的示例使用两个量化后端 QuantoConfig 和 transformers.BitsAndBytesConfig,分别用于 transformer 和文本编码器。
import torch
from diffusers import DiffusionPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers.quantization_config import QuantoConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
pipeline_quant_config = PipelineQuantizationConfig(
quant_mapping={
"transformer": QuantoConfig(weights_dtype="int8"),
"text_encoder_2": TransformersBitsAndBytesConfig(
load_in_4bit=True, compute_dtype=torch.bfloat16
),
}
)Transformers 中有一个独立的 bitsandbytes 后端。对于来自 Transformers 的组件,您需要导入并使用 transformers.BitsAndBytesConfig。例如,FluxPipeline 中的 text_encoder_2 是 Transformers 中的 T5EncoderModel,因此您需要使用 transformers.BitsAndBytesConfig 而不是 diffusers.BitsAndBytesConfig。
如果您不想管理这些不同的导入,或者不确定每个流水线组件的来源,请使用上述的基础量化方法。
import torch
from diffusers import DiffusionPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
pipeline_quant_config = PipelineQuantizationConfig(
quant_mapping={
"transformer": DiffusersBitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16),
"text_encoder_2": TransformersBitsAndBytesConfig(
load_in_4bit=True, compute_dtype=torch.bfloat16
),
}
)将 pipeline_quant_config 传递给 from_pretrained() 以量化该 pipeline。
pipe = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
quantization_config=pipeline_quant_config,
dtype=torch.bfloat16,
).to("cuda") # or "mps", "xpu", "cpu"
image = pipe("photo of a cute dog").images[0]加载预量化检查点
某些量化后端支持加载已经量化并上传到 Hub 的检查点。在这种情况下,量化配置存储在模型的 config.json 中并会被自动读取——您无需创建 PipelineQuantizationConfig 或传入 quantization_config 参数。
import torch
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
"rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4",
dtype=torch.bfloat16,
).to("cuda") # or "mps", "xpu", "cpu"以下后端开箱即支持加载预量化检查点。
| 后端 | 备注 |
|---|---|
| bitsandbytes | Config 保存在 config.json 中,无需额外参数。支持即时量化以及加载预量化检查点。 |
| torchao | 同上。 |
| GGUF | 使用 from_single_file 与 Model 类;不支持 pipeline 级别的加载。 |
| BFL | 使用 from_single_file 加载 Black Forest Labs 的 fp8r 和 nvfp4 单文件检查点;不支持流水线级别的加载。 |
| AutoRound | 仅支持加载;请先使用 AutoRound CLI 或 Python API 进行量化。 |
| Nunchaku Lite | Config 保存在 config.json 中;需要 kernels 包。仅支持加载。 |
| ModelOpt | 仅支持加载预量化模型。 |
| SDNQ | 支持在运行时进行量化以及加载预量化模型;需要 sdnq 包。 |
资源
请查看以下资源,以了解更多有关量化的信息。
如果您是量化新手,我们推荐查看以下与 DeepLearning.AI 合作推出的入门友好课程。
如果您有兴趣添加新的量化方法,请参阅贡献新量化方法指南。
Transformers 量化概述介绍了不同量化后端的优缺点。
阅读探索 Diffusers 中的量化后端博客文章,了解每个量化后端的简要介绍、如何选择合适的后端,以及如何将量化与其他内存优化技术结合使用。