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
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,
torch_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,
torch_dtype=torch.float16,
)默认情况下，所有其他模块（例如 torch.nn.LayerNorm）都会被转换为 torch.float16。您可以使用 torch_dtype 参数更改这些模块的数据类型。
transformer_8bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
+ torch_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,
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 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)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 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,
torch_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,
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 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,
torch_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,
torch_dtype=torch.float16,
)反量化 bitsandbytes 模型
量化后，您可以将模型反量化回其原始精度，但这可能会导致轻微的质量损失。请确保您有足够的 GPU RAM 来容纳反量化后的模型。
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,
torch_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,
torch_dtype=torch.float16,
)
text_encoder_2_4bit.dequantize()
transformer_4bit.dequantize()torch.compile
使用 torch.compile 加速推理。请确保您安装了最新的 bitsandbytes，我们还建议安装 PyTorch nightly。
torch._dynamo.config.capture_dynamic_output_shape_ops = True
quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
torch_dtype=torch.float16,
)
transformer_4bit.compile(fullgraph=True)在 RTX 4090 上，使用编译时，4 位 Flux 生成在 25.809 秒内完成，而未使用编译时为 32.570 秒。
查看基准测试脚本了解更多详情。