bitsandbytes
bitsandbytes 是将模型量化为 8 位和 4 位的最简单选项。8 位量化将 fp16 中的异常值与 int8 中的非异常值相乘,将非异常值转换回 fp16,然后将它们相加以返回 fp16 格式的权重。这减少了异常值对模型性能的负面影响。
4 位量化进一步压缩了模型,通常与 QLoRA 一起使用,以微调量化的大语言模型。
要使用 bitsandbytes,请确保你已安装以下库:
pip install diffusers transformers accelerate bitsandbytes -U
现在你可以通过将一个 [BitsAndBytesConfig
] 传递给 [~ModelMixin.from_pretrained
] 来量化模型。这适用于任何模态的任何模型,只要它支持使用 Accelerate 加载并且包含 torch.nn.Linear
层。
使用 get_memory_footprint
方法检查你的内存占用:
print(model.get_memory_footprint())
量化模型可以从 [~ModelMixin.from_pretrained
] 方法加载,而无需指定 quantization_config
参数:
from diffusers import FluxTransformer2DModel, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_4bit=True)
model_4bit = FluxTransformer2DModel.from_pretrained(
"hf-internal-testing/flux.1-dev-nf4-pkg", subfolder="transformer"
)
8 位 (LLM.int8() 算法)
本节探讨了8位模型的一些特定特性,如异常值阈值和跳过模块转换。
异常值阈值
“异常值”是指超过某个阈值的隐藏状态值,这些值在 fp16 中计算。虽然这些值通常呈正态分布([-3.5, 3.5]),但对于大型模型,这种分布可能会非常不同([-60, 6] 或 [6, 60])。8位量化对于值约为5的情况效果很好,但超过这个范围后,性能会显著下降。一个良好的默认阈值是6,但对于更不稳定的模型(小模型或微调)可能需要更低的阈值。
为了找到适合你模型的最佳阈值,我们建议尝试在 [BitsAndBytesConfig
] 中的 llm_int8_threshold
参数。
from diffusers import FluxTransformer2DModel, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(
load_in_8bit=True, llm_int8_threshold=10,
)
model_8bit = FluxTransformer2DModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quantization_config,
)
跳过模块转换
对于某些模型,你不需要将每个模块都量化到8位,这实际上可能会导致不稳定。例如,对于像 Stable Diffusion 3 这样的扩散模型,可以使用 llm_int8_skip_modules
参数在 [BitsAndBytesConfig
] 中跳过 proj_out
模块。
from diffusers import SD3Transformer2DModel, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(
load_in_8bit=True, llm_int8_skip_modules=["proj_out"],
)
model_8bit = SD3Transformer2DModel.from_pretrained(
"stabilityai/stable-diffusion-3-medium-diffusers",
subfolder="transformer",
quantization_config=quantization_config,
)
4 位 (QLoRA 算法)
本节探讨了 4 位模型的一些特定功能,例如更改计算数据类型、使用 Normal Float 4 (NF4) 数据类型以及使用嵌套量化。
计算数据类型
为了加速计算,你可以使用 BitsAndBytesConfig
中的 bnb_4bit_compute_dtype
参数将数据类型从 float32(默认值)更改为 bf16:
import torch
from diffusers import BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
Normal Float 4 (NF4)
NF4 是一种 4 位数据类型,来自 QLoRA 论文,适用于从正态分布初始化的权重。你应该使用 NF4 来训练 4 位基础模型。这可以通过 BitsAndBytesConfig
中的 bnb_4bit_quant_type
参数进行配置:
from diffusers import BitsAndBytesConfig
nf4_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
)
model_nf4 = SD3Transformer2DModel.from_pretrained(
"stabilityai/stable-diffusion-3-medium-diffusers",
subfolder="transformer",
quantization_config=nf4_config,
)
对于推理,bnb_4bit_quant_type
对性能的影响不大。然而,为了与模型权重保持一致,你应该使用 bnb_4bit_compute_dtype
和 torch_dtype
的值。
嵌套量化
嵌套量化是一种可以在不增加性能成本的情况下节省额外内存的技术。此功能对已经量化的权重进行第二次量化,以节省额外的 0.4 位/参数。
from diffusers import BitsAndBytesConfig
double_quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
)
double_quant_model = SD3Transformer2DModel.from_pretrained(
"stabilityai/stable-diffusion-3-medium-diffusers",
subfolder="transformer",
quantization_config=double_quant_config,
)
反量化 bitsandbytes
模型
一旦量化后,你可以将模型反量化回原始精度,但这可能会导致模型质量有轻微的损失。确保你有足够的 GPU 内存来容纳反量化的模型。
from diffusers import BitsAndBytesConfig
double_quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
)
double_quant_model = SD3Transformer2DModel.from_pretrained(
"stabilityai/stable-diffusion-3-medium-diffusers",
subfolder="transformer",
quantization_config=double_quant_config,
)
model.dequantize()