Diffusers 文档
量化
并获得增强的文档体验
开始使用
量化
量化侧重于用更少的比特表示数据,同时尽量保持原始数据的精度。这通常意味着将数据类型转换为用更少的比特表示相同信息。例如,如果您的模型权重存储为32位浮点数并量化为16位浮点数,这将使模型大小减半,从而更容易存储并减少内存使用。较低的精度还可以加快推理速度,因为用更少的比特执行计算所需的时间更少。
Diffusers 支持多种量化后端,以使大型扩散模型(如 Flux)更易于访问。本指南将展示如何使用 PipelineQuantizationConfig 类在从预训练或未量化检查点初始化管道时对其进行量化。
管道级量化
您可以通过两种方式使用 PipelineQuantizationConfig,具体取决于您对管道中每个模型的量化规格所需控制的级别。
- 对于更基本和简单的用例,您只需定义
quant_backend
、quant_kwargs
和components_to_quantize
- 对于更细粒度的量化控制,提供一个
quant_mapping
,其中包含各个模型组件的量化规格
简单量化
使用以下参数初始化 PipelineQuantizationConfig。
quant_backend
指定要使用的量化后端。目前支持的后端包括:bitsandbytes_4bit
、bitsandbytes_8bit
、gguf
、quanto
和torchao
。quant_kwargs
包含要使用的特定量化参数。components_to_quantize
指定要量化管道的哪些组件。通常,您应该量化计算最密集的组件,例如 Transformer。如果管道具有多个文本编码器(例如 FluxPipeline),则文本编码器是另一个需要考虑量化的组件。下面的示例量化了 FluxPipeline 中的 T5 文本编码器,同时保持 CLIP 模型不变。
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() 以量化管道。
pipe = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
quantization_config=pipeline_quant_config,
torch_dtype=torch.bfloat16,
).to("cuda")
image = pipe("photo of a cute dog").images[0]
quant_mapping
quant_mapping
参数提供了更灵活的选项,用于量化管道中的每个独立组件,例如组合不同的量化后端。
初始化 PipelineQuantizationConfig 并向其传递 quant_mapping
。quant_mapping
允许您指定管道中每个组件的量化选项,例如 Transformer 和文本编码器。
以下示例将 ~quantizers.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() 以量化管道。
pipe = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
quantization_config=pipeline_quant_config,
torch_dtype=torch.bfloat16,
).to("cuda")
image = pipe("photo of a cute dog").images[0]
资源
请查看以下资源以了解更多关于量化的信息。
如果您是量化新手,我们建议您查看以下与 DeepLearning.AI 合作的入门级课程。
如果您有兴趣添加新的量化方法,请参阅贡献新量化方法指南。
Transformers 量化概述提供了不同量化后端的优缺点概览。
阅读探索 Diffusers 中的量化后端博客文章,简要介绍每个量化后端、如何选择后端以及如何将量化与其他内存优化结合使用。