Skip to content

bitsandbytes

bitsandbytes 是将模型量化为 8 位和 4 位的最简单选项。8 位量化将 fp16 中的异常值与 int8 中的非异常值相乘,将非异常值转换回 fp16,然后将它们相加以返回 fp16 格式的权重。这减少了异常值对模型性能的负面影响。

4 位量化进一步压缩了模型,通常与 QLoRA 一起使用,以微调量化的大语言模型。

要使用 bitsandbytes,请确保你已安装以下库:

bash
pip install diffusers transformers accelerate bitsandbytes -U

现在你可以通过将一个 [BitsAndBytesConfig] 传递给 [~ModelMixin.from_pretrained] 来量化模型。这适用于任何模态的任何模型,只要它支持使用 Accelerate 加载并且包含 torch.nn.Linear 层。

使用 get_memory_footprint 方法检查你的内存占用:

py
print(model.get_memory_footprint())

量化模型可以从 [~ModelMixin.from_pretrained] 方法加载,而无需指定 quantization_config 参数:

py
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 参数。

py
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 模块。

py
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:

py
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 参数进行配置:

py
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_dtypetorch_dtype 的值。

嵌套量化

嵌套量化是一种可以在不增加性能成本的情况下节省额外内存的技术。此功能对已经量化的权重进行第二次量化,以节省额外的 0.4 位/参数。

py
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 内存来容纳反量化的模型。

python
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()

资源