Hugging Face 的 LOGO

Diffusers 文件

bitsandbytes

Hugging Face's logo
加入 Hugging Face 社群

並獲得增強的文件體驗

開始使用

bitsandbytes

bitsandbytes 是將模型量化到 8 位和 4 位最簡單的選擇。8 位量化將 fp16 中的離群值與 int8 中的非離群值相乘，將非離群值轉換回 fp16，然後將它們相加，以 fp16 格式返回權重。這減少了離群值對模型效能的負面影響。

4 位量化進一步壓縮模型，常用於 QLoRA 以微調量化的 LLM。

本指南演示了量化如何能使 FLUX.1-dev 在少於 16GB 的 VRAM 上執行，甚至可以在免費的 Google Colab 例項上執行。

comparison image

要使用 bitsandbytes，請確保您已安裝以下庫：

pip install diffusers transformers accelerate bitsandbytes -U

現在，您可以透過將 BitsAndBytesConfig 傳遞給 from_pretrained() 來量化模型。這適用於任何模態的任何模型，只要它支援使用 Accelerate 載入幷包含 torch.nn.Linear 層即可。

8位
4位

將模型量化為 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。

8位
4位
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 秒。

檢視基準測試指令碼瞭解更多詳情。

資源

< > 在 GitHub 上更新

←入門 gguf→
© . This site is unofficial and not affiliated with Hugging Face, Inc.