feat: pipeline-level quantization config

[sayakpaul]

What does this PR do?

See: #10327

TL;DR: This PR adds support to apply a quantization config when doing DiffusionPipeline.from_pretrained(...), thereby making it easier for the users to benefit from quantization.

Why

To apply quantization to a DiffusionPipeline, a user has to first initialize the models they want to quantize with desired quantization_configs:

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 = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    torch_dtype=torch.float16,
)

pipe = DiffusionPipeline.from_pretrained(
    ..., transformer=transformer_8bit, text_encoder_2=text_encoder_2_8bit
)

This is cumbersome.

What

@SunMarc and I worked on this PR to show the kind of simple changes we need to enable a user to pass a quantization config directly whilst doing DiffusionPipeline.from_pretrained(..., quantization_config=...). The user experience now becomes:

pipeline_quant_config = PipelineQuantizationConfig(
    quant_backend="bitsandbytes_8bit",
    quant_kwargs={"load_in_8bit": True},
    modules_to_quantize=["text_encoder_2", "transformer"]
)
pipe = DiffusionPipeline.from_pretrained(
    ...,
    quantization_config=pipeline_quant_config
)

Users can specify granular level quantization mapping too:

quant_config = {
    "transformer": DiffBitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
     ),
    "text_encoder_2": TranBitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
     ),
}

This is particularly helpful when using different quantization backends for different modules (below we show a combination of Quanto and BitsAndBytes):

quant_config = {
    "transformer": QuantoConfig(weights_dtype="float8"),
    "text_encoder_2": TranBitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
     ),
}

Here's a script that might be helpful for others to test:

code

from diffusers.quantizers import PipelineQuantizationConfig
from diffusers import DiffusionPipeline
import argparse
import torch


def get_global_config():
    quant_config = PipelineQuantizationConfig(
        quant_backend="bitsandbytes_4bit",
        quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
        modules_to_quantize=["transformer", "text_encoder_2"],
    )
    return quant_config


def get_granular_config(use_quanto=False):
    from diffusers import BitsAndBytesConfig as DiffBitsAndBytesConfig, QuantoConfig
    from transformers import BitsAndBytesConfig as TranBitsAndBytesConfig

    transformer_config = (
        QuantoConfig(weights_dtype="float8")
        if use_quanto
        else DiffBitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
        )
    )

    quant_config = {
        "transformer": transformer_config,
        "text_encoder_2": TranBitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
        ),
    }
    return quant_config


def load_pipeline(quant_config):
    pipe = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev", quantization_config=quant_config, torch_dtype=torch.bfloat16
    ).to("cuda")
    return pipe


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--use_global_config", action="store_true")
    parser.add_argument("--use_quanto", action="store_true")
    args = parser.parse_args()

    quant_config = get_global_config() if args.use_global_config else get_granular_config(args.use_quanto)
    pipe = load_pipeline(quant_config)

    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]
    image.save(f"quant_global@{args.use_global_config}_quanto@{args.use_quanto}.png")

Cc: @asomoza if you want to test this out :)

TODOs

• Docs
• Tests

[sayakpaul]

Should there be a reasonable default for this? @SunMarc had some ideas around this.

[SunMarc]

I was thinking it could be nice to have a class attribute e.g modules_to_quantize in each pipeline. Or we can just create a mapping pipeline <-> modules_to_quantize if you prefer to keep this outside of the class. (e.g just like how peft deal with modules_to_target for loras)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

I think that we can go ahead and start adding simple test and a bit of documentation !