bitsandbytes

bitsandbytes 是将模型量化为 8 位和 4 位的最简单方案。8 位量化将 fp16 中的离群值与 int8 中的非离群值相乘，将非离群值转换回 fp16，然后将两者相加，从而以 fp16 格式返回权重。这降低了离群值对模型性能产生的负面影响。

4 位量化可以进一步压缩模型，并且通常与 QLoRA 结合使用，以微调量化后的大语言模型。

本指南演示了如何通过量化，在不到 16GB 显存上运行 FLUX.1-dev，甚至可以在免费的 Google Colab 实例上运行。

comparison image

要使用 bitsandbytes，请确保已安装以下库：

pip install diffusers transformers accelerate bitsandbytes -U

现在，您可以通过向 from_pretrained() 传入 BitsAndBytesConfig 来量化模型。只要模型支持使用 Accelerate 加载，并且包含 torch.nn.Linear 层，此方法就适用于任何模态下的任何模型。

8 位

4 位

将模型量化为 8 位可将内存占用减半：

bitsandbytes 在 Transformers 和 Diffusers 中均受支持，因此您可以对 FluxTransformer2DModel 和 T5EncoderModel 进行量化。

对于 Ada 及更高系列的 GPU，我们建议将 dtype 更改为 torch.bfloat16。

CLIPTextModel 和 AutoencoderKL 未进行量化，因为它们本身体积已经很小，而且 AutoencoderKL 仅包含少数几个 torch.nn.Linear 层。

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
import torch
from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(load_in_8bit=True,)

text_encoder_2_8bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True,)

transformer_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

默认情况下，torch.nn.LayerNorm 等其他所有模块都会被转换为 torch.float16。您可以通过 dtype 参数更改这些模块的数据类型。

transformer_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
+   dtype=torch.float32,
)

让我们使用量化后的模型生成一张图像。

设置 device_map="auto" 会优先填满所有可用的 GPU 空间，然后是 CPU，最后，如果内存仍然不足，则使用硬盘（速度最慢的选项）。

from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=transformer_8bit,
    text_encoder_2=text_encoder_2_8bit,
    dtype=torch.float16,
    device_map="auto",
)

pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 50,
    "max_sequence_length": 512,
}

image = pipe(**pipe_kwargs, generator=torch.manual_seed(0),).images[0]

当内存充足时，您也可以使用 .to("cuda") 直接将 pipeline 移至 GPU，并应用 enable_model_cpu_offload() 来优化 GPU 内存使用。

Once a model is quantized, you can push the model to the Hub with the push_to_hub() method. The quantization config.json file is pushed first, followed by the quantized model weights. You can also save the serialized 8-bit models locally with save_pretrained().

将模型量化为 4 位可将内存占用降低 4 倍：

bitsandbytes 在 Transformers 和 Diffusers 中均受支持，因此您可以对 FluxTransformer2DModel 和 T5EncoderModel 进行量化。

对于 Ada 及更高系列的 GPU，我们建议将 dtype 更改为 torch.bfloat16。

CLIPTextModel 和 AutoencoderKL 未进行量化，因为它们本身体积已经很小，而且 AutoencoderKL 仅包含少数几个 torch.nn.Linear 层。

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
import torch
from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(load_in_4bit=True,)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(load_in_4bit=True,)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

默认情况下，torch.nn.LayerNorm 等其他所有模块都会被转换为 torch.float16。您可以通过 dtype 参数更改这些模块的数据类型。

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
+   dtype=torch.float32,
)

让我们使用量化后的模型生成一张图像。

设置 device_map="auto" 会优先自动填满所有可用的 GPU 空间，然后是 CPU，最后，如果内存仍然不足，则使用硬盘（绝对最慢的选项）。

from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    transformer=transformer_4bit,
    text_encoder_2=text_encoder_2_4bit,
    dtype=torch.float16,
    device_map="auto",
)

pipe_kwargs = {
    "prompt": "A cat holding a sign that says hello world",
    "height": 1024,
    "width": 1024,
    "guidance_scale": 3.5,
    "num_inference_steps": 50,
    "max_sequence_length": 512,
}

image = pipe(**pipe_kwargs, generator=torch.manual_seed(0),).images[0]

当内存充足时，您也可以使用 .to("cuda") 直接将 pipeline 移至 GPU，并应用 enable_model_cpu_offload() 来优化 GPU 内存使用。

模型量化完成后，您可以使用 push_to_hub() 方法将模型推送到 Hub。量化 config.json 文件会先被推送，随后是量化后的模型权重。您也可以通过 save_pretrained() 将序列化后的 4-bit 模型保存到本地。

使用 8 位和 4 位权重进行训练仅支持训练额外参数。

使用 get_memory_footprint 方法检查您的内存占用：

print(model.get_memory_footprint())

请注意，这仅能告知您模型参数的内存占用，而不会估算推理所需的内存。

量化模型可以直接从 from_pretrained() 方法加载，无需指定 quantization_config 参数：

from diffusers import AutoModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_4bit=True)

model_4bit = AutoModel.from_pretrained(
    "hf-internal-testing/flux.1-dev-nf4-pkg", subfolder="transformer"
)

8 位（LLM.int8() 算法）

有关 8 位量化的详细信息，请参阅这篇博客文章！

本节将探讨 8 位模型的一些特定功能，例如异常值阈值和跳过模块转换。

异常值阈值

“离群值”是指大于特定阈值的隐藏状态值，这些值以 fp16 计算。虽然这些值通常呈正态分布（[-3.5, 3.5]），但对于大型模型，其分布可能差异很大（[-60, 6] 或 [6, 60]）。8 位量化在数值约为 5 时效果良好，但超出该范围后，性能会显著下降。一个好的默认阈值是 6，但对于更不稳定的模型（小型模型或微调），可能需要更低的阈值。

要找到适合您模型的最佳阈值，我们建议尝试调整 BitsAndBytesConfig 中的 llm_int8_threshold 参数：

from diffusers import AutoModel, BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True, llm_int8_threshold=10,
)

model_8bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quantization_config,
)

跳过模块转换

对于某些模型，您无需将所有模块都量化为 8 位，因为这实际上可能会导致不稳定性。例如，对于 Stable Diffusion 3 等扩散模型，可以使用 BitsAndBytesConfig 中的 llm_int8_skip_modules 参数来跳过 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)

常规 Float 4 (NF4)

NF4 是 QLoRA 论文中提出的一种 4 位数据类型，专为从正态分布初始化的权重而设计。在训练 4 位基础模型时，应使用 NF4。可以通过 BitsAndBytesConfig 中的 bnb_4bit_quant_type 参数进行配置：

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

在推理阶段，bnb_4bit_quant_type 对性能的影响并不大。但为了与模型权重保持一致，您应使用 bnb_4bit_compute_dtype 和 dtype 的值。

嵌套量化

嵌套量化是一种可以在不增加额外性能开销的情况下进一步节省内存的技术。该功能会对已量化的权重进行第二次量化，从而额外节省 0.4 位/参数。

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

反量化 bitsandbytes 模型

量化完成后，您可以将模型反量化回其原始精度，但这可能会导致轻微的质量损失。请确保您的 GPU 显存足以容纳反量化后的模型。

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

from diffusers import AutoModel
from transformers import T5EncoderModel

quant_config = TransformersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

text_encoder_2_4bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
)

transformer_4bit = AutoModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    dtype=torch.float16,
)

text_encoder_2_4bit.dequantize()
transformer_4bit.dequantize()

torch.compile

使用 torch.compile 加速推理。请确保已安装最新版本的 bitsandbytes，同时我们建议安装 PyTorch nightly。

8 位

4 位

在配备编译功能的 RTX 4090 上，4 位 Flux 生成耗时 25.809 秒，而未启用编译时耗时 32.570 秒。

有关更多详细信息，请查看基准测试脚本。

资源

Update on GitHub