Quanto

Quanto 后端已弃用，并将在 1.0.0 版本中移除。请考虑切换到其他受支持的量化后端，例如 bitsandbytes 或 torchao。

Quanto 是 Optimum 的 PyTorch 量化后端。它的设计兼顾了通用性和简洁性：

所有功能在 eager 模式下均可用（适用于不可追踪的模型）
支持量化感知训练
量化模型与 torch.compile 兼容
量化模型与设备无关（例如 CUDA、XPU、MPS、CPU）

要使用 Quanto 后端，您首先需要安装 optimum-quanto>=0.2.6 和 accelerate

pip install optimum-quanto accelerate

现在，您可以通过将 QuantoConfig 对象传递给 from_pretrained() 方法来量化模型。尽管 Quanto 库允许对 nn.Conv2d 和 nn.LayerNorm 模块进行量化，但目前 Diffusers 仅支持量化模型中 nn.Linear 层的权重。以下代码片段演示了如何使用 Quanto 应用 float8 量化。

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")

在特定模块上跳过量化

可以使用 QuantoConfig 中的 modules_to_not_convert 参数来跳过对某些模块应用量化。请确保传入该参数的模块与 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,
)

使用 from_single_file 与 Quanto 后端

QuantoConfig 与 ~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)

保存量化模型

Diffusers 支持使用 ~ModelMixin.save_pretrained 方法序列化 Quanto 模型。

使用 Quanto 库直接量化的模型与使用 Diffusers（以 Quanto 作为后端）量化的模型，其序列化和加载要求有所不同。目前无法使用 ~ModelMixin.from_pretrained 将直接用 Quanto 量化的模型加载到 Diffusers 中。

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>")

使用 torch.compile 与 Quanto

目前，Quanto 后端支持以下量化类型的 torch.compile：

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")

支持的量化类型

权重

float8
int8
int4
int2

Update on GitHub