Diffusers 文档
bitsandbytes
并获得增强的文档体验
开始
bitsandbytes
bitsandbytes 是将模型量化为 8 位和 4 位的最简单选项。8 位量化将 fp16 中的异常值与 int8 中的非异常值相乘，将非异常值转换回 fp16，然后将它们加在一起以返回 fp16 中的权重。这减少了异常值对模型性能的降级影响。
4 位量化进一步压缩模型，并且通常与 QLoRA 一起用于微调量化的 LLM。
本指南演示了量化如何使 FLUX.1-dev 能够在小于 16GB VRAM 的环境中运行，甚至可以在免费的 Google Colab 实例上运行。
要使用 bitsandbytes，请确保您已安装以下库
pip install diffusers transformers accelerate bitsandbytes -U
现在，您可以通过将 BitsAndBytesConfig 传递给 from_pretrained() 来量化模型。这适用于任何模态的任何模型，只要它支持使用 Accelerate 加载并包含 torch.nn.Linear 层。
以 8 位量化模型可将内存使用量减半
Transformers 和 Diffusers 都支持 bitsandbytes，因此您可以量化 FluxTransformer2DModel 和 T5EncoderModel。
对于 Ada 及更高系列的 GPU，我们建议将 torch_dtype 更改为 torch.bfloat16。
CLIPTextModel 和 AutoencoderKL 未被量化，因为它们本身尺寸已经很小，并且 AutoencoderKL 只有少量的 torch.nn.Linear 层。
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
from diffusers import FluxTransformer2DModel
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,
torch_dtype=torch.float16,
)
quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True,)
transformer_8bit = FluxTransformer2DModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)默认情况下，所有其他模块（如 torch.nn.LayerNorm）都将转换为 torch.float16。您可以使用 torch_dtype 参数更改这些模块的数据类型。
transformer_8bit = FluxTransformer2DModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
+ torch_dtype=torch.float32,
)让我们使用量化模型生成图像。
设置 device_map="auto" 会自动首先填充 GPU 上的所有可用空间，然后是 CPU，最后，如果仍然没有足够的内存，则填充硬盘驱动器（绝对最慢的选项）。
pipe = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=transformer_8bit,
text_encoder_2=text_encoder_2_8bit,
torch_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() 在本地保存序列化的 8 位模型。
仅支持使用 8 位和 4 位权重训练额外参数。
使用 get_memory_footprint 方法检查您的内存占用
print(model.get_memory_footprint())可以从 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 位量化的更多详细信息！
本节探讨了 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 这样的扩散模型，可以使用 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)Normal Float 4 (NF4)
NF4 是 QLoRA 论文中的 4 位数据类型，适用于从正态分布初始化的权重。您应该使用 NF4 训练 4 位基础模型。这可以使用 BitsAndBytesConfig 中的 bnb_4bit_quant_type 参数进行配置
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
from diffusers import FluxTransformer2DModel
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,
torch_dtype=torch.float16,
)
quant_config = DiffusersBitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
)
transformer_4bit = FluxTransformer2DModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)对于推理，bnb_4bit_quant_type 对性能没有巨大影响。但是，为了与模型权重保持一致，您应该使用 bnb_4bit_compute_dtype 和 torch_dtype 值。
嵌套量化
嵌套量化是一种可以节省额外内存且不增加性能成本的技术。此功能对已量化的权重执行二次量化，以额外节省 0.4 比特/参数。
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
from diffusers import FluxTransformer2DModel
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,
torch_dtype=torch.float16,
)
quant_config = DiffusersBitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
)
transformer_4bit = FluxTransformer2DModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)解量化 bitsandbytes 模型
一旦量化，您可以将模型解量化到其原始精度，但这可能会导致少量质量损失。请确保您有足够的 GPU RAM 来容纳解量化后的模型。
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
from diffusers import FluxTransformer2DModel
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,
torch_dtype=torch.float16,
)
quant_config = DiffusersBitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
)
transformer_4bit = FluxTransformer2DModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
text_encoder_2_4bit.dequantize()
transformer_4bit.dequantize()