Support full int8 quantization for diffusers (#1193)

* Support full int8 quantization for diffusers

* Add SD quant API test

* apply black

* apply style

* reapply black

* Apply suggestions from code review

---------

Co-authored-by: Ilyas Moutawwakil <[email protected]>

Author: l-bat

optimum/commands/export/openvino.py

@@ -364,6 +364,9 @@ def run(self):
                     )
 
                 if self.args.quant_mode in ["nf4_f8e4m3", "nf4_f8e5m2", "int4_f8e4m3", "int4_f8e5m2"]:
+                    if library_name == "diffusers":
+                        raise NotImplementedError("Mixed precision quantization isn't supported for diffusers.")
+
                     wc_config = prepare_wc_config(self.args, _DEFAULT_4BIT_CONFIG)
                     wc_dtype, q_dtype = self.args.quant_mode.split("_")
                     wc_config["dtype"] = wc_dtype
@@ -421,7 +424,7 @@ def run(self):
 
                 model_cls = OVSanaPipeline
             else:
-                raise NotImplementedError(f"Quantization in hybrid mode isn't supported for class {class_name}.")
+                raise NotImplementedError(f"Quantization isn't supported for class {class_name}.")
 
             model = model_cls.from_pretrained(self.args.model, export=True, quantization_config=quantization_config)
             model.save_pretrained(self.args.output)

optimum/intel/openvino/configuration.py

@@ -789,10 +789,12 @@ def post_init(self):
 
         if self.dataset is not None:
             speech_to_text_datasets = list(PREDEFINED_SPEECH_TO_TEXT_DATASETS.keys())
-            if self.dataset not in LANGUAGE_DATASETS + speech_to_text_datasets:
+            stable_diffusion_datasets = list(PREDEFINED_SD_DATASETS.keys())
+            if self.dataset not in LANGUAGE_DATASETS + speech_to_text_datasets + stable_diffusion_datasets:
                 raise ValueError(
-                    f"""You can only choose between the following datasets: {LANGUAGE_DATASETS} for LLMs or
-                    {speech_to_text_datasets} for speech-to-text models, but we found {self.dataset}."""
+                    f"""You can only choose between the following datasets: {LANGUAGE_DATASETS} for LLMs,
+                    {speech_to_text_datasets} for speech-to-text models or
+                    {stable_diffusion_datasets} for diffusion models, but we found {self.dataset}."""
                 )
 
         if self.bits != 8:

optimum/intel/openvino/modeling_diffusion.py

@@ -534,7 +534,7 @@ def _from_pretrained(
         else:
             # why is this quantization not performed in __init__?
             if ov_pipeline_class.export_feature != "text-to-image":
-                raise NotImplementedError(f"Quantization in hybrid mode is not supported for {cls.__name__}")
+                raise NotImplementedError(f"Quantization is not supported for {cls.__name__}")
 
             from optimum.intel import OVQuantizer
 
@@ -548,10 +548,13 @@ def _from_pretrained(
             # same as in DiffusionPipeline.from_pretrained, we save where the model was instantiated from
             ov_pipeline.register_to_config(_name_or_path=config.get("_name_or_path", str(model_id)))
 
-            hybrid_quantization_config = deepcopy(quantization_config)
-            hybrid_quantization_config.quant_method = OVQuantizationMethod.HYBRID
             quantizer = OVQuantizer(ov_pipeline)
-            quantizer.quantize(ov_config=OVConfig(quantization_config=hybrid_quantization_config))
+            if isinstance(quantization_config, OVWeightQuantizationConfig):
+                hybrid_quantization_config = deepcopy(quantization_config)
+                hybrid_quantization_config.quant_method = OVQuantizationMethod.HYBRID
+                quantizer.quantize(ov_config=OVConfig(quantization_config=hybrid_quantization_config))
+            else:
+                quantizer.quantize(ov_config=OVConfig(quantization_config=quantization_config))
 
             return ov_pipeline
         ov_pipeline = ov_pipeline_class(
@@ -878,9 +881,11 @@ def reshape(
             self.text_encoder.model = self._reshape_text_encoder(
                 self.text_encoder.model,
                 batch_size,
-                getattr(self.tokenizer, "model_max_length", -1)
-                if "Gemma" not in self.tokenizer.__class__.__name__
-                else -1,
+                (
+                    getattr(self.tokenizer, "model_max_length", -1)
+                    if "Gemma" not in self.tokenizer.__class__.__name__
+                    else -1
+                ),
             )
 
         if self.text_encoder_2 is not None:

optimum/intel/openvino/quantization.py

@@ -440,6 +440,16 @@ def _quantize_ovbasemodel(
             # Quantize model(s)
             if isinstance(self.model, _OVModelForWhisper):
                 self._quantize_whisper_model(quantization_config, calibration_dataset, **kwargs)
+            elif is_diffusers_available() and isinstance(self.model, OVDiffusionPipeline):
+                for name, sub_model in self.model.ov_submodels.items():
+                    if name not in ("unet", "transformer"):
+                        _weight_only_quantization(sub_model, OVWeightQuantizationConfig(bits=8), **kwargs)
+                    else:
+                        quantized_vision_model = _full_quantization(
+                            sub_model, quantization_config, calibration_dataset, **kwargs
+                        )
+                        getattr(self.model, name).model = quantized_vision_model
+                self.model.clear_requests()
             else:
                 quantized_model = _full_quantization(
                     self.model.model, quantization_config, calibration_dataset, **kwargs
@@ -450,6 +460,9 @@ def _quantize_ovbasemodel(
             if calibration_dataset is None:
                 raise ValueError("Calibration dataset is required to run quantization.")
 
+            if is_diffusers_available() and isinstance(self.model, OVDiffusionPipeline):
+                raise NotImplementedError("Mixed precision quantization isn't supported for diffusers.")
+
             quantized_model = _mixed_quantization(self.model.model, quantization_config, calibration_dataset, **kwargs)
             self.model.model = quantized_model
             self.model.request = None

tests/openvino/test_exporters_cli.py

@@ -189,19 +189,23 @@ class OVCLIExportTestCase(unittest.TestCase):
             "int8",
             "--dataset librispeech --num-samples 1 --smooth-quant-alpha 0.9 --trust-remote-code",
             [14, 22, 21] if is_transformers_version("<=", "4.36.0") else [14, 22, 25],
-            [{"int8": 14}, {"int8": 21}, {"int8": 17}]
-            if is_transformers_version("<=", "4.36.0")
-            else [{"int8": 14}, {"int8": 22}, {"int8": 18}],
+            (
+                [{"int8": 14}, {"int8": 21}, {"int8": 17}]
+                if is_transformers_version("<=", "4.36.0")
+                else [{"int8": 14}, {"int8": 22}, {"int8": 18}]
+            ),
         ),
         (
             "automatic-speech-recognition-with-past",
             "whisper",
             "f8e4m3",
             "--dataset librispeech --num-samples 1 --smooth-quant-alpha 0.9 --trust-remote-code",
             [14, 22, 21] if is_transformers_version("<=", "4.36.0") else [14, 22, 25],
-            [{"f8e4m3": 14}, {"f8e4m3": 21}, {"f8e4m3": 17}]
-            if is_transformers_version("<=", "4.36.0")
-            else [{"f8e4m3": 14}, {"f8e4m3": 22}, {"f8e4m3": 18}],
+            (
+                [{"f8e4m3": 14}, {"f8e4m3": 21}, {"f8e4m3": 17}]
+                if is_transformers_version("<=", "4.36.0")
+                else [{"f8e4m3": 14}, {"f8e4m3": 22}, {"f8e4m3": 18}]
+            ),
         ),
         (
             "text-generation",
@@ -263,6 +267,62 @@ class OVCLIExportTestCase(unittest.TestCase):
                 {"f8e5m2": 2, "int4": 28},
             ],
         ),
+        (
+            "stable-diffusion",
+            "stable-diffusion",
+            "int8",
+            "--dataset conceptual_captions --num-samples 1 --trust-remote-code",
+            [
+                112,
+                0,
+                0,
+                0,
+            ],
+            [
+                {"int8": 121},
+                {"int8": 42},
+                {"int8": 34},
+                {"int8": 64},
+            ],
+        ),
+        (
+            "stable-diffusion-xl",
+            "stable-diffusion-xl",
+            "f8e5m2",
+            "--dataset laion/220k-GPT4Vision-captions-from-LIVIS --num-samples 1 --trust-remote-code",
+            [
+                174,
+                0,
+                0,
+                0,
+                0,
+            ],
+            [
+                {"f8e5m2": 183},
+                {"int8": 42},
+                {"int8": 34},
+                {"int8": 64},
+                {"int8": 66},
+            ],
+        ),
+        (
+            "latent-consistency",
+            "latent-consistency",
+            "f8e4m3",
+            "--dataset laion/filtered-wit --num-samples 1 --trust-remote-code",
+            [
+                79,
+                0,
+                0,
+                0,
+            ],
+            [
+                {"f8e4m3": 84},
+                {"int8": 42},
+                {"int8": 34},
+                {"int8": 40},
+            ],
+        ),
     ]
 
     TEST_4BIT_CONFIGURATIONS = [
@@ -709,6 +769,8 @@ def test_exporters_cli_full_quantization(
                     expected_fake_nodes_per_model = expected_fake_nodes_per_model[:-1]
             elif "text-generation" in task:
                 submodels = [model]
+            elif any(x in task for x in ("stable-diffusion", "latent-consistency")):
+                submodels = model.ov_submodels.values()
             else:
                 raise Exception("Unexpected task.")
 

tests/openvino/test_quantization.py

@@ -19,26 +19,22 @@
 import itertools
 import logging
 import unittest
-from collections import defaultdict, Iterable
+from collections import defaultdict
+from collections.abc import Iterable
 from enum import Enum
 from functools import partial
 from typing import Union, Type
 
 import openvino as ov
 import pytest
-import evaluate
 import numpy as np
 import torch
-from datasets import load_dataset
 from parameterized import parameterized
 import nncf
 from transformers import (
     AutoModelForQuestionAnswering,
-    AutoModelForSequenceClassification,
     AutoTokenizer,
     AutoProcessor,
-    TrainingArguments,
-    default_data_collator,
 )
 from transformers.testing_utils import slow
 from transformers.utils.quantization_config import QuantizationMethod
@@ -116,9 +112,11 @@ class OVQuantizerTest(unittest.TestCase):
                 smooth_quant_alpha=0.95,
             ),
             [14, 22, 21] if is_transformers_version("<=", "4.36.0") else [14, 22, 25],
-            [{"int8": 14}, {"int8": 21}, {"int8": 17}]
-            if is_transformers_version("<=", "4.36.0")
-            else [{"int8": 14}, {"int8": 22}, {"int8": 18}],
+            (
+                [{"int8": 14}, {"int8": 21}, {"int8": 17}]
+                if is_transformers_version("<=", "4.36.0")
+                else [{"int8": 14}, {"int8": 22}, {"int8": 18}]
+            ),
         ),
         (
             OVModelForCausalLM,
@@ -234,6 +232,77 @@ class OVQuantizerTest(unittest.TestCase):
                 {"f8e5m2": 2, "int4": 28},
             ],
         ),
+        (
+            OVStableDiffusionPipeline,
+            "stable-diffusion",
+            dict(
+                weight_only=False,
+                dataset="conceptual_captions",
+                num_samples=1,
+                processor=MODEL_NAMES["stable-diffusion"],
+                trust_remote_code=True,
+            ),
+            [
+                112,
+                0,
+                0,
+                0,
+            ],
+            [
+                {"int8": 121},
+                {"int8": 42},
+                {"int8": 34},
+                {"int8": 64},
+            ],
+        ),
+        (
+            OVStableDiffusionXLPipeline,
+            "stable-diffusion-xl",
+            dict(
+                weight_only=False,
+                dtype="f8e5m2",
+                dataset="laion/220k-GPT4Vision-captions-from-LIVIS",
+                num_samples=1,
+                processor=MODEL_NAMES["stable-diffusion-xl"],
+                trust_remote_code=True,
+            ),
+            [
+                174,
+                0,
+                0,
+                0,
+                0,
+            ],
+            [
+                {"f8e5m2": 183},
+                {"int8": 42},
+                {"int8": 34},
+                {"int8": 64},
+                {"int8": 66},
+            ],
+        ),
+        (
+            OVLatentConsistencyModelPipeline,
+            "latent-consistency",
+            OVQuantizationConfig(
+                dtype="f8e4m3",
+                dataset="laion/filtered-wit",
+                num_samples=1,
+                trust_remote_code=True,
+            ),
+            [
+                79,
+                0,
+                0,
+                0,
+            ],
+            [
+                {"f8e4m3": 84},
+                {"int8": 42},
+                {"int8": 34},
+                {"int8": 40},
+            ],
+        ),
     ]
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES_TORCH_MODEL)
@@ -359,6 +428,11 @@ def test_ov_model_static_quantization_with_auto_dataset(
                 tokens = tokenizer("This is a sample input", return_tensors="pt")
                 outputs = ov_model(**tokens)
                 self.assertTrue("logits" in outputs)
+            elif any(
+                x == model_cls
+                for x in (OVStableDiffusionPipeline, OVStableDiffusionXLPipeline, OVLatentConsistencyModelPipeline)
+            ):
+                submodels = ov_model.ov_submodels.values()
             else:
                 raise Exception("Unexpected model class.")