Diffusers documentation
bitsandbytes
bitsandbytes
bitsandbytes là lựa chọn dễ nhất để lượng tử hóa mô hình xuống 8 và 4-bit. Lượng tử hóa 8-bit nhân các giá trị ngoại lai trong fp16 với các giá trị không ngoại lai trong int8, chuyển các giá trị không ngoại lai trở lại fp16, sau đó cộng chúng lại để trả về trọng số ở fp16. Điều này làm giảm tác động suy giảm mà các giá trị ngoại lai gây ra đối với hiệu suất của mô hình.
Lượng tử hóa 4-bit nén mô hình sâu hơn nữa và thường được sử dụng cùng với QLoRA để tinh chỉnh các LLM đã được lượng tử hóa.
Hướng dẫn này cho thấy cách lượng tử hóa có thể cho phép chạy FLUX.1-dev trên dưới 16GB VRAM và thậm chí trên một instance Google Colab miễn phí.
Để sử dụng bitsandbytes, hãy đảm bảo bạn đã cài đặt các thư viện sau:
pip install diffusers transformers accelerate bitsandbytes -U
Bây giờ bạn có thể lượng tử hóa một mô hình bằng cách truyền BitsAndBytesConfig vào from_pretrained(). Cách này hoạt động với bất kỳ mô hình nào trong bất kỳ chế độ nào, miễn là nó hỗ trợ tải bằng Accelerate và chứa các lớp torch.nn.Linear.
Việc lượng tử hóa một mô hình ở 8-bit sẽ giảm một nửa dung lượng bộ nhớ sử dụng:
bitsandbytes được hỗ trợ trong cả Transformers và Diffusers, vì vậy bạn có thể lượng tử hóa cả FluxTransformer2DModel và T5EncoderModel.
Đối với các GPU thuộc dòng Ada và cao hơn, chúng tôi khuyến nghị thay đổi dtype thành torch.bfloat16.
CLIPTextModelvà AutoencoderKL không được lượng tử hóa vì chúng đã có kích thước nhỏ và vì AutoencoderKL chỉ có một vài lớptorch.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,
)Mặc định, tất cả các module khác như torch.nn.LayerNorm sẽ được chuyển đổi sang torch.float16. Bạn có thể thay đổi kiểu dữ liệu của các module này bằng tham số dtype.
transformer_8bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
+ dtype=torch.float32,
)Hãy tạo một hình ảnh bằng cách sử dụng các mô hình đã lượng tử hóa của chúng ta.
Việc đặt device_map="auto" sẽ tự động lấp đầy toàn bộ dung lượng trống trên GPU trước, sau đó là CPU, và cuối cùng là ổ cứng (tùy chọn chậm nhất tuyệt đối) nếu vẫn còn thiếu bộ nhớ.
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]
Khi có đủ bộ nhớ, bạn cũng có thể trực tiếp di chuyển pipeline sang GPU bằng .to("cuda") và áp dụng enable_model_cpu_offload() để tối ưu hóa việc sử dụng bộ nhớ GPU.
Sau khi một model được lượng tử hóa, bạn có thể đẩy model lên Hub bằng phương thức push_to_hub(). File cấu hình lượng tử hóa config.json sẽ được đẩy lên trước, tiếp theo là trọng số model đã lượng tử hóa. Bạn cũng có thể lưu các model 8-bit đã được serialize cục bộ bằng save_pretrained().
Việc lượng tử hóa mô hình ở 4-bit giúp giảm 4 lần lượng bộ nhớ sử dụng:
bitsandbytes được hỗ trợ trong cả Transformers và Diffusers, vì vậy bạn có thể lượng tử hóa cả FluxTransformer2DModel và T5EncoderModel.
Đối với các GPU thuộc dòng Ada và cao hơn, chúng tôi khuyến nghị thay đổi dtype thành torch.bfloat16.
CLIPTextModelvà AutoencoderKL không được lượng tử hóa vì chúng đã có kích thước nhỏ và vì AutoencoderKL chỉ có một vài lớptorch.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,
)Mặc định, tất cả các module khác như torch.nn.LayerNorm sẽ được chuyển đổi sang torch.float16. Bạn có thể thay đổi kiểu dữ liệu của các module này bằng tham số dtype.
transformer_4bit = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
quantization_config=quant_config,
+ dtype=torch.float32,
)Hãy tạo một hình ảnh bằng cách sử dụng các mô hình đã lượng tử hóa của chúng ta.
Việc đặt device_map="auto" sẽ tự động lấp đầy toàn bộ dung lượng trống trên GPU trước, sau đó là CPU, và cuối cùng là ổ cứng (tùy chọn chậm nhất tuyệt đối) nếu vẫn còn thiếu bộ nhớ.
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]
Khi có đủ bộ nhớ, bạn cũng có thể trực tiếp di chuyển pipeline sang GPU bằng .to("cuda") và áp dụng enable_model_cpu_offload() để tối ưu hóa việc sử dụng bộ nhớ GPU.
Sau khi một model được lượng tử hóa, bạn có thể đẩy model lên Hub bằng phương thức push_to_hub(). File cấu hình lượng tử hóa config.json sẽ được đẩy lên trước, tiếp theo là trọng số model đã lượng tử hóa. Bạn cũng có thể lưu các model 4-bit đã được serialize cục bộ bằng save_pretrained().
Huấn luyện với trọng số 8-bit và 4-bit chỉ được hỗ trợ cho việc huấn luyện các tham số bổ sung.
Kiểm tra mức sử dụng bộ nhớ của bạn bằng phương thức get_memory_footprint:
print(model.get_memory_footprint())Lưu ý rằng điều này chỉ cho bạn biết dung lượng bộ nhớ của các tham số mô hình và không ước tính yêu cầu bộ nhớ cho suy luận.
Các mô hình đã lượng tử hóa có thể được tải từ phương thức from_pretrained() mà không cần chỉ định các tham số 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-bit (thuật toán LLM.int8())
Tìm hiểu thêm về các chi tiết của lượng tử hóa 8-bit trong bài đăng trên blog này!
Phần này khám phá một số tính năng cụ thể của các mô hình 8-bit, chẳng hạn như ngưỡng ngoại lệ và việc bỏ qua chuyển đổi module.
Ngưỡng ngoại lai
Một “outlier” là giá trị trạng thái ẩn lớn hơn một ngưỡng nhất định, và các giá trị này được tính toán ở fp16. Mặc dù các giá trị thường có phân phối chuẩn ([-3.5, 3.5]), nhưng phân phối này có thể rất khác nhau đối với các mô hình lớn ([-60, 6] hoặc [6, 60]). Lượng tử hóa 8-bit hoạt động tốt cho các giá trị ~5, nhưng vượt quá mức đó, sẽ có sự suy giảm hiệu suất đáng kể. Một giá trị ngưỡng mặc định tốt là 6, nhưng có thể cần một ngưỡng thấp hơn cho các mô hình kém ổn định hơn (mô hình nhỏ hoặc finetuning).
Để tìm ngưỡng tối ưu cho mô hình của bạn, chúng tôi khuyến nghị thử nghiệm với tham số llm_int8_threshold trong BitsAndBytesConfig:
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,
)Bỏ qua chuyển đổi module
Đối với một số mô hình, bạn không cần phải lượng tử hóa mọi module xuống 8-bit vì điều này thực sự có thể gây ra sự bất ổn. Ví dụ, đối với các mô hình khuếch tán như Stable Diffusion 3, module proj_out có thể được bỏ qua bằng cách sử dụng tham số llm_int8_skip_modules trong BitsAndBytesConfig:
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-bit (thuật toán QLoRA)
Tìm hiểu thêm về các chi tiết của nó trong bài đăng trên blog này.
Phần này khám phá một số tính năng cụ thể của các mô hình 4-bit, chẳng hạn như thay đổi kiểu dữ liệu tính toán, sử dụng kiểu dữ liệu Normal Float 4 (NF4) và sử dụng lượng tử hóa lồng nhau.
Loại dữ liệu tính toán
Để tăng tốc độ tính toán, bạn có thể thay đổi kiểu dữ liệu từ float32 (giá trị mặc định) sang bf16 bằng cách sử dụng tham số bnb_4bit_compute_dtype trong BitsAndBytesConfig:
import torch
from diffusers import BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)Float 4 thông thường (NF4)
NF4 là kiểu dữ liệu 4-bit từ bài báo QLoRA, được điều chỉnh cho các trọng số được khởi tạo từ phân phối chuẩn. Bạn nên sử dụng NF4 để huấn luyện các mô hình cơ bản 4-bit. Điều này có thể được cấu hình bằng tham số bnb_4bit_quant_type trong BitsAndBytesConfig:
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,
)Đối với suy luận, bnb_4bit_quant_type không có tác động lớn đến hiệu suất. Tuy nhiên, để nhất quán với trọng số của mô hình, bạn nên sử dụng các giá trị bnb_4bit_compute_dtype và dtype.
Lượng tử hóa lồng nhau
Lượng tử hóa lồng nhau là một kỹ thuật có thể tiết kiệm thêm bộ nhớ mà không làm tăng chi phí hiệu năng. Tính năng này thực hiện lượng tử hóa lần hai đối với các trọng số đã được lượng tử hóa để tiết kiệm thêm 0,4 bit/thông số.
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,
)Khử lượng hóa các mô hình bitsandbytes
Sau khi lượng tử hóa, bạn có thể khử lượng tử hóa mô hình về độ chính xác ban đầu, nhưng điều này có thể dẫn đến một mức độ suy giảm chất lượng nhỏ. Hãy đảm bảo bạn có đủ bộ nhớ RAM GPU để chứa mô hình đã khử lượng tử hóa.
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
Tăng tốc suy luận bằng torch.compile. Hãy đảm bảo bạn đã cài đặt phiên bản bitsandbytes mới nhất và chúng tôi cũng khuyến nghị cài đặt 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,
dtype=torch.float16,
)
transformer_4bit.compile(fullgraph=True)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,
)
transformer_4bit.compile(fullgraph=True)Trên RTX 4090 với biên dịch, việc tạo ảnh Flux 4-bit hoàn thành trong 25.809 giây so với 32.570 giây khi không có.
Xem script benchmarking để biết thêm chi tiết.