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},
    components_to_quantize=["text_encoder_2", "transformer"]
)
pipe = DiffusionPipeline.from_pretrained(
    ...,
    quantization_config=pipeline_quant_config
)

Users can specify granular level quantization mapping too:

quant_mapping = {
    "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
     ),
}
pipeline_quant_config = PipelineQuantizationConfig(quant_mapping=quant_mapping)

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

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

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},
        components_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_mapping = {
        "transformer": transformer_config,
        "text_encoder_2": TranBitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16
        ),
    }
    quant_config = PipelineQuantizationConfig(quant_mapping=quant_mapping)
    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
• Enable testing in CI

[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 !

[sayakpaul]

@DN6 it would be great if you could do a first review of this

[SunMarc]

cc @DerekLiu35 might be interesting to you !

[yiyixuxu]

the change to pipelines code looks good to me!

the provided code example isn't working though (does not throw an error but no effect) - do we intend to support passing as a dict?

[DN6]

Looking good. Could we add a test as well please.

[sayakpaul]

@stevhliu where do you think we should document it?

[stevhliu]

We can document the two ways to pass quantization configs directly on the overview or we can create a new doc and add it after the overview.

I think adding directly to the overview would probably be the easiest.

It would also be nice to recommend when to use each method. With this new PipelineQuantizationConfig (which seems much simpler), are there still scenarios where a user would want to initialize the model first and then pass a quantization config?

[sayakpaul]

We can document the two ways to pass quantization configs directly on the overview or we can create a new doc and add it after the overview.

I think adding directly to the overview would probably be the easiest.

I feel the same. Done in 872c91e.

With this new PipelineQuantizationConfig (which seems much simpler), are there still scenarios where a user would want to initialize the model first and then pass a quantization config?

Good question. I don't think there would be any need. We will see in due time.

[yiyixuxu]

@sayakpaul
I'm not suggesting passing as a dict, I made that comment because in the code you provided for test in the PR description, you passed it as a dict and the pipeline does not throw an error

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

[sayakpaul]

@yiyixuxu my bad. I have updated the code snippet. I have also added a check for validating the quantization_config input to DiffusionPipeline (and a test for it).

[yiyixuxu]

thanks!

[SunMarc]

Thanks a lot ! LGTM !

[DN6]

Minor comments that are not merge blocks. Can be addressed here or in a follow up.

[sayakpaul]

Thanks all for your comments! Cc: @asomoza @apolinario @linoytsaban for awareness of this feature as it makes the barrier to entry for quantization of a DiffusionPipeline lots easier.