Add hybrid quantization for StableDiffusion pipelines

Author: l-bat

optimum/intel/openvino/configuration.py

@@ -178,7 +178,8 @@ class OVWeightQuantizationConfig(QuantizationConfigMixin):
                     using the [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
         dataset (`Union[List[str]]`, *optional*):
             The dataset used for data-aware compression. You can provide your own dataset in a list of string or just use the
-            the one from the list ['wikitext2','c4','c4-new','ptb','ptb-new']
+            the one from the list ['wikitext2','c4','c4-new','ptb','ptb-new'] for LLLMs or
+            ['conceptual_captions','laion/220k-GPT4Vision-captions-from-LIVIS','laion/filtered-wit'] for SD models
         group_size (`int`, *optional*, defaults to 128):
             The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization.
         ratio (`float`, *optional*, defaults to 1.0):
@@ -193,6 +194,8 @@ class OVWeightQuantizationConfig(QuantizationConfigMixin):
             Enables AWQ method to unify weight ranges and improve overall model accuracy.
         ignored_scope (`nncf.IgnoredScope`, *optional*):
             An ignored scope that defined the list of model control flow graph nodes to be ignored during quantization.
+        subset_size (`int`, *optional*, defaults to 128):
+            Number of data samples to calculate activation statistics.
 
     """
 
@@ -207,6 +210,7 @@ def __init__(
         all_layers: Optional[bool] = None,
         sensitivity_metric: Optional[str] = None,
         ignored_scope: Optional[dict] = None,
+        subset_size: int = 128,
         **kwargs,
     ):
         self.bits = bits
@@ -218,6 +222,7 @@ def __init__(
         self.all_layers = all_layers
         self.sensitivity_metric = sensitivity_metric
         self.ignored_scope = ignored_scope
+        self.subset_size = subset_size
         self.quant_method = "default"  # TODO : enable AWQ after nncf v2.9.0 release
         self.post_init()
 
@@ -230,10 +235,16 @@ def post_init(self):
         if self.group_size is not None and self.group_size != -1 and self.group_size <= 0:
             raise ValueError("`group_size` must be greater than 0 or equal to -1")
         if self.dataset is not None and isinstance(self.dataset, str):
-            if self.dataset not in ["wikitext2", "c4", "c4-new", "ptb", "ptb-new"]:
+            llm_datasets = ["wikitext2", "c4", "c4-new", "ptb", "ptb-new"]
+            stable_diffusion_datasets = [
+                "conceptual_captions",
+                "laion/220k-GPT4Vision-captions-from-LIVIS",
+                "laion/filtered-wit"
+            ]
+            if self.dataset not in llm_datasets + stable_diffusion_datasets:
                 raise ValueError(
                     f"""You have entered a string value for dataset. You can only choose between
-                    ['wikitext2','c4','c4-new','ptb','ptb-new'], but we found {self.dataset}"""
+                    {llm_datasets} for LLLMs or {stable_diffusion_datasets} for SD models, but we found {self.dataset}"""
                 )
 
         if self.bits not in [4, 8]:

optimum/intel/openvino/modeling_diffusion.py

@@ -14,6 +14,7 @@
 
 import importlib
 import logging
+import math
 import os
 import shutil
 from pathlib import Path
@@ -261,9 +262,19 @@ def _from_pretrained(
         if load_in_8bit:
             quantization_config = quantization_config or {"bits": 8}
 
-        unet = cls.load_model(
-            new_model_save_dir / DIFFUSION_MODEL_UNET_SUBFOLDER / unet_file_name, quantization_config
-        )
+        if isinstance(quantization_config, dict):
+            quantization_config = OVWeightQuantizationConfig.from_dict(quantization_config)
+
+        dataset = None
+        if quantization_config:
+            dataset = quantization_config.dataset
+            quantization_config.dataset = None  # apply weight compression without dataset
+
+        unet_path = new_model_save_dir / DIFFUSION_MODEL_UNET_SUBFOLDER / unet_file_name
+        if quantization_config and dataset is None:
+            unet = cls.load_model(unet_path, quantization_config)
+        else:
+            unet = cls.load_model(unet_path)
 
         components = {
             "vae_encoder": new_model_save_dir / DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER / vae_encoder_file_name,
@@ -278,8 +289,58 @@ def _from_pretrained(
         if model_save_dir is None:
             model_save_dir = new_model_save_dir
 
+        if quantization_config and dataset is not None:
+            sd_model = cls(unet=unet, config=config, model_save_dir=model_save_dir, **components, **kwargs)
+
+            supported_pipelines = (
+                OVStableDiffusionPipeline,
+                OVStableDiffusionXLPipeline,
+                OVLatentConsistencyModelPipeline
+            )
+            if not isinstance(sd_model, supported_pipelines):
+                raise NotImplementedError(f"Quantization in hybrid mode is not supported for {cls.__name__}")
+
+            num_inference_steps = 4 if isinstance(cls, OVLatentConsistencyModelPipeline) else 50
+            quantization_config.dataset = dataset
+
+            if isinstance(quantization_config.dataset, str):
+                from .quantization import get_stable_diffusion_dataset
+                dataset_name = quantization_config.dataset
+                num_samples = math.ceil(quantization_config.subset_size / num_inference_steps)
+                quantization_config.dataset = get_stable_diffusion_dataset(dataset_name, num_samples)
+
+            unet_inputs = sd_model.prepare_inputs(quantization_config.dataset, quantization_config.subset_size, num_inference_steps)
+            quantization_config.dataset = unet_inputs
+
+            from .quantization import _hybrid_quantization
+            unet = _hybrid_quantization(sd_model.unet.model, quantization_config)
+
         return cls(unet=unet, config=config, model_save_dir=model_save_dir, **components, **kwargs)
 
+    def prepare_inputs(
+        self,
+        dataset: "Dataset",
+        subset_size: int,
+        num_inference_steps: int,
+        height: Optional[int] = 512,
+        width: Optional[int] = 512,
+        **kwargs,
+    ) -> "Dataset":
+        self.compile()
+        calibration_data = []
+
+        from .quantization import InferRequestWrapper
+        self.unet.request = InferRequestWrapper(self.unet.request, calibration_data)
+        for prompt in dataset.get_inference_data():
+            _ = self.__call__(prompt, num_inference_steps=num_inference_steps, height=height, width=width)
+            if len(calibration_data) >= subset_size:
+                break
+        self.unet.request = self.unet.request.request
+
+        from nncf import Dataset
+        return Dataset(calibration_data)
+
+
     @classmethod
     def _from_transformers(
         cls,

optimum/intel/openvino/quantization.py

@@ -15,6 +15,9 @@
 import inspect
 import logging
 import os
+from collections import deque
+from copy import deepcopy
+from datasets import load_dataset
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
 
@@ -23,6 +26,7 @@
 import torch
 import transformers
 from nncf import CompressWeightsMode, IgnoredScope, NNCFConfig, SensitivityMetric
+from nncf.quantization.advanced_parameters import AdvancedSmoothQuantParameters
 from nncf.torch import create_compressed_model, register_default_init_args, register_module
 from nncf.torch.dynamic_graph.io_handling import wrap_nncf_model_inputs_with_objwalk
 from nncf.torch.initialization import PTInitializingDataLoader
@@ -577,4 +581,104 @@ def _weight_only_quantization(
         # awq=config.quant_method == "awq", # TODO : remove and add it back once nncf v2.9.0
         ignored_scope=ignored_scope,
         dataset=dataset,
+        subset_size=config.subset_size,
     )
+
+
+def _get_operation_const_op(operation, const_port_id: int):
+    node = operation.input_value(const_port_id).get_node()
+    queue = deque([node])
+    constant_node = None
+    allowed_propagation_types_list = ["Convert", "FakeQuantize", "Reshape"]
+
+    while len(queue) != 0:
+        curr_node = queue.popleft()
+        if curr_node.get_type_name() == "Constant":
+            constant_node = curr_node
+            break
+        if len(curr_node.inputs()) == 0:
+            break
+        if curr_node.get_type_name() in allowed_propagation_types_list:
+            queue.append(curr_node.input_value(0).get_node())
+
+    return constant_node
+
+
+def _is_embedding(node) -> bool:
+    allowed_types_list = ["f16", "f32", "f64"]
+    const_port_id = 0
+    input_tensor = node.input_value(const_port_id)
+    if input_tensor.get_element_type().get_type_name() in allowed_types_list:
+        const_node = _get_operation_const_op(node, const_port_id)
+        if const_node is not None:
+            return True
+
+    return False
+
+
+def _collect_ops_with_weights(model):
+    ops_with_weights = []
+    for op in model.get_ops():
+        if op.get_type_name() == "MatMul":
+            constant_node_0 = _get_operation_const_op(op, const_port_id=0)
+            constant_node_1 = _get_operation_const_op(op, const_port_id=1)
+            if constant_node_0 or constant_node_1:
+                ops_with_weights.append(op.get_friendly_name())
+        if op.get_type_name() == "Gather" and _is_embedding(op):
+            ops_with_weights.append(op.get_friendly_name())
+
+    return ops_with_weights
+
+
+def get_stable_diffusion_dataset(
+    dataset_name: str, nsamples: int = 50, seed: int = 0, text_column: str = "caption"
+) -> nncf.Dataset:
+    if dataset_name not in [
+            "conceptual_captions",
+            "laion/220k-GPT4Vision-captions-from-LIVIS",
+            "laion/filtered-wit"
+        ]:
+        raise ValueError(
+            f"""You have entered a string value for dataset. You can only choose between
+             ['conceptual_captions','laion/220k-GPT4Vision-captions-from-LIVIS','laion/filtered-wit'],
+             but we found {dataset_name}"""
+        )
+
+    data = load_dataset(dataset_name, split="train", streaming=True).shuffle(seed=seed).take(nsamples)
+    dataset = [batch[text_column] for batch in data]
+    return nncf.Dataset(dataset)
+
+
+def _hybrid_quantization(
+    model: openvino.runtime.Model, quantization_config: Union[OVWeightQuantizationConfig, Dict]
+):
+    dataset = quantization_config.dataset
+    wc_ignored_scope = deepcopy(quantization_config.ignored_scope)
+
+    if isinstance(wc_ignored_scope, dict):
+        wc_ignored_scope["types"] = wc_ignored_scope.get("types", []) + ["Convolution"]
+    else:
+        assert wc_ignored_scope is None
+        wc_ignored_scope = {"types": ["Convolution"]}
+
+    ops_to_compress = _collect_ops_with_weights(model)
+    ptq_ignored_scope = deepcopy(quantization_config.ignored_scope)
+    if isinstance(ptq_ignored_scope, dict):
+        ptq_ignored_scope["names"] = ptq_ignored_scope.get("names", []) + ops_to_compress
+    else:
+        assert ptq_ignored_scope is None
+        ptq_ignored_scope = {"names": ops_to_compress}
+
+    quantization_config.dataset = None  # Apply Weight Compression without dataset
+    quantization_config.ignored_scope = wc_ignored_scope
+    compressed_model = _weight_only_quantization(model, quantization_config)
+
+    quantized_model = nncf.quantize(
+        compressed_model,
+        dataset,
+        model_type=nncf.ModelType.TRANSFORMER,
+        ignored_scope=nncf.IgnoredScope(**ptq_ignored_scope),
+        advanced_parameters=nncf.AdvancedQuantizationParameters(AdvancedSmoothQuantParameters(matmul=-1)),
+        subset_size=quantization_config.subset_size,
+    )
+    return quantized_model

tests/openvino/test_quantization.py

@@ -36,6 +36,7 @@
 
 from optimum.intel import (
     OVConfig,
+    OVLatentConsistencyModelPipeline,
     OVModelForAudioClassification,
     OVModelForCausalLM,
     OVModelForFeatureExtraction,
@@ -230,6 +231,12 @@ class OVWeightCompressionTest(unittest.TestCase):
         (OVStableDiffusionXLPipeline, "stable-diffusion-xl"),
     )
 
+    SUPPORTED_ARCHITECTURES_WITH_HYBRID_QUANTIZATION = (
+        (OVStableDiffusionPipeline, "stable-diffusion", 72, 195),
+        (OVStableDiffusionXLPipeline, "stable-diffusion-xl", 84, 331),
+        (OVLatentConsistencyModelPipeline, "latent-consistency", 50, 135),
+    )
+
     IS_SUPPORT_STATEFUL = is_openvino_version(">=", "2023.3")
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_EXPECTED_8BIT_COMPRESSED_MATMULS)
@@ -345,6 +352,40 @@ def test_ovmodel_load_with_compressed_weights(self, model_cls, model_type):
             _, num_int8, _ = get_num_quantized_nodes(model)
             self.assertEqual(expected_ov_int8[i], num_int8)
 
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_HYBRID_QUANTIZATION)
+    def test_ovmodel_hybrid_quantization(self, model_cls, model_type, expected_num_fake_quantize, expected_ov_int8):
+        model_id = MODEL_NAMES[model_type]
+        quantization_config = OVWeightQuantizationConfig(bits=8, dataset="conceptual_captions", subset_size=5)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            model = model_cls.from_pretrained(model_id, export=True, quantization_config=quantization_config)
+
+            num_fake_quantize, num_int8, num_int4 = get_num_quantized_nodes(model.unet)
+            self.assertEqual(expected_num_fake_quantize, num_fake_quantize)
+            self.assertEqual(expected_ov_int8, num_int8)
+            self.assertEqual(0, num_int4)
+
+            model.save_pretrained(tmp_dir)
+
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_HYBRID_QUANTIZATION)
+    def test_ovmodel_hybrid_quantization_with_custom_dataset(
+        self, model_cls, model_type, expected_num_fake_quantize, expected_ov_int8
+    ):
+        model_id = MODEL_NAMES[model_type]
+        dataset_name = "daspartho/stable-diffusion-prompts"
+        dataset = load_dataset(dataset_name, split="train", streaming=True)
+        quantization_dataset = nncf.Dataset(dataset, lambda x: x["prompt"])
+        model = model_cls.from_pretrained(
+            model_id,
+            export=True,
+            quantization_config=OVWeightQuantizationConfig(
+                bits=8, dataset=quantization_dataset, subset_size=3
+            ),
+        )
+        num_fake_quantize, num_int8, num_int4 = get_num_quantized_nodes(model.unet)
+        self.assertEqual(expected_num_fake_quantize, num_fake_quantize)
+        self.assertEqual(expected_ov_int8, num_int8)
+        self.assertEqual(0, num_int4)
+
     @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_EXPECTED_4BIT_AUTOCOMPRESSED_MATMULS)
     def test_ovmodel_4bit_auto_compression(self, model_cls, model_type, expected_ov_int8, expected_ov_int4):
         with tempfile.TemporaryDirectory() as tmp_dir:

tests/openvino/utils_tests.py

@@ -127,8 +127,8 @@ def get_num_quantized_nodes(ov_model):
         if "FakeQuantize" in elem.name:
             num_fake_quantize += 1
         for i in range(elem.get_output_size()):
-            if "8" in elem.get_output_element_type(i).get_type_name():
+            if elem.get_output_element_type(i).get_type_name() in ["i8", "u8"]:
                 num_int8 += 1
-            if "4" in elem.get_output_element_type(i).get_type_name():
+            if elem.get_output_element_type(i).get_type_name() in ["i4", "u4"]:
                 num_int4 += 1
     return num_fake_quantize, num_int8, num_int4