Quanto

Hậu端 Quanto đã bị ngừng hỗ trợ và sẽ bị loại bỏ trong phiên bản 1.0.0. Hãy cân nhắc chuyển sang một trong các hậu端 lượng tử hóa khác được hỗ trợ, chẳng hạn như bitsandbytes hoặc torchao.

Quanto là backend lượng tử hóa PyTorch cho Optimum. Nó được thiết kế với sự linh hoạt và đơn giản:

Tất cả các tính năng đều khả dụng ở chế độ eager (hoạt động với các mô hình không thể trace)
Hỗ trợ huấn luyện nhận thức lượng tử hóa
Các mô hình lượng tử hóa tương thích với torch.compile
Các mô hình đã lượng tử hóa không phụ thuộc vào thiết bị (ví dụ CUDA, XPU, MPS, CPU)

Để sử dụng backend Quanto, trước tiên bạn cần cài đặt optimum-quanto>=0.2.6 và accelerate

pip install optimum-quanto accelerate

Bây giờ bạn có thể lượng tử hóa một mô hình bằng cách truyền đối tượng QuantoConfig vào phương thức from_pretrained(). Mặc dù thư viện Quanto cho phép lượng tử hóa các module nn.Conv2d và nn.LayerNorm, hiện tại Diffusers chỉ hỗ trợ lượng tử hóa trọng số trong các lớp nn.Linear của một mô hình. Đoạn mã sau minh họa cách áp dụng lượng tử hóa float8 với Quanto.

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)

pipe = FluxPipeline.from_pretrained(model_id, transformer=transformer, dtype=dtype)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A cat holding a sign that says hello world"
image = pipe(
    prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512
).images[0]
image.save("output.png")

Bỏ qua Lượng tử hóa trên các module cụ thể

Bạn có thể bỏ qua việc áp dụng lượng tử hóa cho một số module nhất định bằng cách sử dụng đối số modules_to_not_convert trong QuantoConfig. Vui lòng đảm bảo rằng các module được truyền vào đối số này khớp với các khóa của các module trong state_dict

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8", modules_to_not_convert=["proj_out"])
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)

Sử dụng from_single_file với Backend Quanto

QuantoConfig tương thích với ~FromOriginalModelMixin.from_single_file.

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_single_file(ckpt_path, quantization_config=quantization_config, dtype=torch.bfloat16)

Lưu trữ các mô hình đã lượng tử hóa

Diffusers hỗ trợ xuất mô hình Quanto bằng phương thức ~ModelMixin.save_pretrained.

Yêu cầu về việc serialize và nạp khác nhau đối với các mô hình được lượng tử hóa trực tiếp bằng thư viện Quanto và các mô hình được lượng tử hóa bằng Diffusers sử dụng Quanto làm backend. Hiện tại, không thể nạp các mô hình được lượng tử hóa trực tiếp bằng Quanto vào Diffusers bằng ~ModelMixin.from_pretrained

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)
# save quantized model to reuse
transformer.save_pretrained("<your quantized model save path>")

# you can reload your quantized model with
model = FluxTransformer2DModel.from_pretrained("<your quantized model save path>")

Sử dụng torch.compile với Quanto

Hiện tại, backend Quanto hỗ trợ torch.compile cho các loại lượng tử hóa sau:

trọng số int8

import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="int8")
transformer = FluxTransformer2DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)

pipe = FluxPipeline.from_pretrained(
    model_id, transformer=transformer, dtype=dtype
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"
images = pipe("A cat holding a sign that says hello").images[0]
images.save("flux-quanto-compile.png")

Các loại lượng hóa được hỗ trợ

Trọng số

float8
int8
int4
int2

Update on GitHub