[WIP] Pipeline level quant config

[sayakpaul]

Just a PoC to see things.

[SunMarc]

Added a few comments !

[sayakpaul]

@SunMarc WDYT about the latest updates?

[SunMarc]

Thanks for iterating ! Much better. I think that we can open an PR on diffusers library now wdyt ?

[sayakpaul]

@SunMarc could you do another pass please? Just want to ensure I got the basic things right?

Maybe any lightweight test you wanna run?

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

[SunMarc]

LGTM overall, just a few nits

[SunMarc]

One thing i'm asking myself is that we should try to provide to the user the best experience. Is it truly the best default to quantize every component in the pipeline or maybe we should pre-define the component should be quantized. WDYT ?
I'm asking this because we will have to add potentially in every snippet the modules_to_quantize as you did in your example. The average users also might not know the name of each component + don't know what component should be quantized or not.

[sayakpaul]

maybe we should pre-define the component should be quantized. WDYT ?

We cannot determine that I guess. For example, users may want to quantize CLIP but that is not typically.

Users can quantize the VAE but that is only supported by a backend that supports quantizing Conv2D (or Conv3D) layers. So, I am not sure what is the best course of action here.

[sayakpaul]

Or maybe we could discuss this in the main PR we'd open to diffusers?

[SunMarc]

We cannot determine that I guess. For example, users may want to quantize CLIP but that is not typically.

We still let them do that if they want but I just want to provide a good default. For example, we can put ['transformers'] by default for some pipeline.

[sayakpaul]

You mean in the docs? Or default to this value in the PipelineQuantConfig class itself?

[sayakpaul]

Specific pipeline docs pages or what are you referring to?

I guess for the docs, we will have a guide in the general overview section of Quantization in diffusers: https://huggingface.co/docs/diffusers/main/en/quantization/overview. There we can do a general recommendation.

Or are you referring to something entirely different?

[SunMarc]

I was thinking about adding a variable to FluxPipeline class potentially for example

[sayakpaul]

Sorry I am not following this at all.

How would the rough changes look like in that case? In my mind, I think users would just want to supply a pipeline-level quant config while doing DiffusionPipeline.from_pretrained() and be done with it.

[sayakpaul]

If we need to add a variable like modules_to_quantize that is a bigger discussion we should be having with the rest of the team. But IMO we are likely ready to take this PR and show how changes for allowing pipeline-level quant configs would look like.

WDYT?

[SunMarc]

yeah, let's open the pr on diffusers repo