rework dataset processing

Author: l-bat

optimum/intel/openvino/configuration.py

@@ -178,9 +178,9 @@ class OVWeightQuantizationConfig(QuantizationConfigMixin):
                 - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved
                     using the [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
         dataset (`str or 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'] for LLLMs or
-            ['conceptual_captions','laion/220k-GPT4Vision-captions-from-LIVIS','laion/filtered-wit'] for SD models.
+            The dataset used for data-aware compression or quantization with NNCF. 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'] for LLLMs
+            or ['conceptual_captions','laion/220k-GPT4Vision-captions-from-LIVIS','laion/filtered-wit'] for SD models.
         ratio (`float`, defaults to 1.0):
             The ratio between baseline and backup precisions (e.g. 0.9 means 90% of layers quantized to INT4_ASYM
             and the rest to INT8_ASYM).

optimum/intel/openvino/modeling_diffusion.py

@@ -14,7 +14,6 @@
 
 import importlib
 import logging
-import math
 import os
 import shutil
 from copy import deepcopy
@@ -59,7 +58,13 @@
 from .configuration import OVConfig, OVWeightQuantizationConfig
 from .loaders import OVTextualInversionLoaderMixin
 from .modeling_base import OVBaseModel
-from .utils import ONNX_WEIGHTS_NAME, OV_TO_NP_TYPE, OV_XML_FILE_NAME, _print_compiled_model_properties
+from .utils import (
+    ONNX_WEIGHTS_NAME,
+    OV_TO_NP_TYPE,
+    OV_XML_FILE_NAME,
+    PREDEFINED_SD_DATASETS,
+    _print_compiled_model_properties,
+)
 
 
 core = Core()
@@ -276,13 +281,15 @@ def _from_pretrained(
                     kwargs[name] = load_method(new_model_save_dir)
 
         quantization_config = cls._prepare_weight_quantization_config(quantization_config, load_in_8bit)
-        weight_quantization_config = deepcopy(quantization_config)
+
+        dataset = None
         unet_path = new_model_save_dir / DIFFUSION_MODEL_UNET_SUBFOLDER / unet_file_name
-        if weight_quantization_config is not None and weight_quantization_config.dataset is not None:
+        if quantization_config is not None and quantization_config.dataset is not None:
+            dataset = quantization_config.dataset
             # load the UNet model uncompressed to apply hybrid quantization further
             unet = cls.load_model(unet_path)
             # Apply weights compression to other `components` without dataset
-            weight_quantization_config.dataset = None
+            quantization_config.dataset = None
         else:
             unet = cls.load_model(unet_path, quantization_config)
 
@@ -294,12 +301,12 @@ def _from_pretrained(
         }
 
         for key, value in components.items():
-            components[key] = cls.load_model(value, weight_quantization_config) if value.is_file() else None
+            components[key] = cls.load_model(value, quantization_config) if value.is_file() else None
 
         if model_save_dir is None:
             model_save_dir = new_model_save_dir
 
-        if quantization_config and quantization_config.dataset is not None:
+        if dataset is not None:
             sd_model = cls(unet=unet, config=config, model_save_dir=model_save_dir, **components, **kwargs)
 
             supported_pipelines = (
@@ -310,24 +317,13 @@ def _from_pretrained(
             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(sd_model, OVLatentConsistencyModelPipeline) else 50
             nsamples = quantization_config.num_samples if quantization_config.num_samples else 200
-            dataset = deepcopy(quantization_config.dataset)
-
-            if isinstance(dataset, str):
-                from .quantization import get_stable_diffusion_dataset
-
-                num_unet_runs = math.ceil(nsamples / num_inference_steps)
-                dataset = get_stable_diffusion_dataset(dataset, num_unet_runs)
-
-            unet_inputs = sd_model._prepare_unet_inputs(dataset, nsamples, num_inference_steps)
+            unet_inputs = sd_model._prepare_unet_inputs(dataset, nsamples)
 
             from .quantization import _hybrid_quantization
 
-            hybrid_quantization_config = deepcopy(quantization_config)
-            hybrid_quantization_config.dataset = unet_inputs
-            hybrid_quantization_config.num_samples = nsamples
-            unet = _hybrid_quantization(sd_model.unet.model, hybrid_quantization_config)
+            unet = _hybrid_quantization(sd_model.unet.model, quantization_config, dataset=unet_inputs)
+            quantization_config.dataset = dataset
 
         return cls(
             unet=unet,
@@ -340,21 +336,52 @@ def _from_pretrained(
 
     def _prepare_unet_inputs(
         self,
-        dataset: List[str],
+        dataset: Union[str, List[Any]],
         num_samples: int,
-        num_inference_steps: int,
         height: Optional[int] = 512,
         width: Optional[int] = 512,
+        seed: Optional[int] = 42,
         **kwargs,
     ) -> Dict[str, Any]:
         self.compile()
-        calibration_data = []
+
+        if isinstance(dataset, str):
+            dataset = deepcopy(dataset)
+            available_datasets = PREDEFINED_SD_DATASETS.keys()
+            if dataset not in available_datasets:
+                raise ValueError(
+                    f"""You have entered a string value for dataset. You can only choose between
+                    {list(available_datasets)}, but the {dataset} was found"""
+                )
+
+            from datasets import load_dataset
+
+            dataset_metadata = PREDEFINED_SD_DATASETS[dataset]
+            dataset = load_dataset(dataset, split=dataset_metadata["split"], streaming=True).shuffle(seed=seed)
+            input_names = dataset_metadata["inputs"]
+            dataset = dataset.select_columns(list(input_names.values()))
+
+            def transform_fn(data_item):
+                return {inp_name: data_item[column] for inp_name, column in input_names.items()}
+
+        else:
+
+            def transform_fn(data_item):
+                return data_item if isinstance(data_item, (list, dict)) else [data_item]
 
         from .quantization import InferRequestWrapper
 
+        calibration_data = []
         self.unet.request = InferRequestWrapper(self.unet.request, calibration_data)
-        for prompt in dataset:
-            _ = self.__call__(prompt, num_inference_steps=num_inference_steps, height=height, width=width)
+
+        for inputs in dataset:
+            inputs = transform_fn(inputs)
+            if isinstance(inputs, dict):
+                self.__call__(**inputs, height=height, width=width)
+            else:
+                self.__call__(*inputs, height=height, width=width)
+            if len(calibration_data) > num_samples:
+                break
 
         self.unet.request = self.unet.request.request
         return calibration_data[:num_samples]

optimum/intel/openvino/quantization.py

@@ -17,15 +17,13 @@
 import logging
 import os
 from collections import deque
-from copy import deepcopy
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Union
 
 import nncf
 import openvino
 import torch
 import transformers
-from datasets import load_dataset
 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
@@ -638,23 +636,8 @@ def _collect_ops_with_weights(model):
     return ops_with_weights
 
 
-def get_stable_diffusion_dataset(
-    dataset_name: str, nsamples: int = 50, seed: int = 0, text_column: str = "caption"
-) -> List[str]:
-    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 dataset
-
-
 def _hybrid_quantization(
-    model: openvino.runtime.Model, quantization_config: OVWeightQuantizationConfig
+    model: openvino.runtime.Model, quantization_config: OVWeightQuantizationConfig, dataset: Dict[str, Any]
 ) -> openvino.runtime.Model:
     """
     Quantize a model in hybrid mode with NNCF which means that we quantize:
@@ -666,28 +649,28 @@ def _hybrid_quantization(
             The OpenVINO Runtime model for applying hybrid quantization.
         quantization_config (`OVWeightQuantizationConfig`):
             The configuration containing the parameters related to quantization.
+        dataset (`Dict[str, Any]`):
+            The dataset used for hybrid quantization.
     Returns:
         The OpenVINO Runtime model with applied hybrid quantization.
     """
-    ignored_scope = quantization_config.ignored_scope if quantization_config.ignored_scope is not None else {}
-
     ops_to_compress = _collect_ops_with_weights(model)
-    ptq_ignored_scope = deepcopy(ignored_scope)
-    ptq_ignored_scope["names"] = ignored_scope.get("names", []) + ops_to_compress
 
-    wc_quantization_config = deepcopy(quantization_config)
+    ignored_scope = quantization_config.ignored_scope if isinstance(quantization_config.ignored_scope, dict) else {}
+    ptq_ignored_scope = nncf.IgnoredScope(**ignored_scope)
+    ptq_ignored_scope.names += ops_to_compress
+
+    wc_quantization_config = copy.deepcopy(quantization_config)
     wc_quantization_config.ignored_scope = ignored_scope
     wc_quantization_config.ignored_scope["types"] = ignored_scope.get("types", []) + ["Convolution"]
-    # Apply Weight Compression without dataset
-    wc_quantization_config.dataset = None
     compressed_model = _weight_only_quantization(model, wc_quantization_config)
 
     subset_size = quantization_config.num_samples if quantization_config.num_samples else 200
     quantized_model = nncf.quantize(
         model=compressed_model,
-        calibration_dataset=nncf.Dataset(quantization_config.dataset),
+        calibration_dataset=nncf.Dataset(dataset),
         model_type=nncf.ModelType.TRANSFORMER,
-        ignored_scope=nncf.IgnoredScope(**ptq_ignored_scope),
+        ignored_scope=ptq_ignored_scope,
         # The SQ algo should be disabled for MatMul nodes because their weights are already compressed
         advanced_parameters=nncf.AdvancedQuantizationParameters(AdvancedSmoothQuantParameters(matmul=-1)),
         subset_size=subset_size,

optimum/intel/openvino/utils.py

@@ -99,6 +99,13 @@
 }
 
 
+PREDEFINED_SD_DATASETS = {
+    "conceptual_captions": {"split": "train", "inputs": {"prompt": "caption"}},
+    "laion/220k-GPT4Vision-captions-from-LIVIS": {"split": "train", "inputs": {"prompt": "caption"}},
+    "laion/filtered-wit": {"split": "train", "inputs": {"prompt": "caption"}},
+}
+
+
 def use_external_data_format(num_parameters: int) -> bool:
     """
     Returns whether or not the model requires using external data format for the ONNX export