Refactor OV weight compression call inside from_pretrained (#683)

* Move calibration dataset construction to WC function

* Tweak tokenizer

* Removed not used import

* ruff

* ruff 2

* Refactor through OVQuantizer call

Author: nikita-savelyevv

optimum/intel/openvino/modeling_decoder.py

@@ -11,7 +11,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
+import copy
 import logging
 import os
 import warnings
@@ -25,7 +25,7 @@
 from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
 from openvino.preprocess import PrePostProcessor
 from openvino.runtime import Core, Tensor, Type
-from transformers import AutoModelForCausalLM, AutoTokenizer, PretrainedConfig
+from transformers import AutoModelForCausalLM, PretrainedConfig
 from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
 from transformers.generation import GenerationMixin
 from transformers.modeling_outputs import CausalLMOutputWithPast
@@ -646,9 +646,8 @@ def _from_pretrained(
                 raise ImportError(
                     "Quantization of the weights requires nncf, please install it with `pip install nncf`"
                 )
-            import nncf
 
-            from .quantization import _weight_only_quantization
+            from optimum.intel.openvino.quantization import OVQuantizer
 
             default_config = _check_default_4bit_configs(config)
 
@@ -657,18 +656,10 @@ def _from_pretrained(
                     f"For the given model, we recommend the following `quantization_config` : {default_config}"
                 )
 
-            calibration_dataset = None
-            if isinstance(quantization_config.dataset, str):
-                tokenizer = quantization_config.tokenizer or AutoTokenizer.from_pretrained(model_id)
-
-                from optimum.gptq.data import get_dataset, prepare_dataset
-
-                nsamples = quantization_config.num_samples or 128
-                dataset = get_dataset(quantization_config.dataset, tokenizer, seqlen=32, nsamples=nsamples)
-                dataset = prepare_dataset(dataset)
-                calibration_dataset = nncf.Dataset(dataset, lambda x: causal_model.prepare_inputs(**x))
-
-            _weight_only_quantization(model, quantization_config, calibration_dataset)
+            quantizer = OVQuantizer(causal_model)
+            quantization_config_copy = copy.deepcopy(quantization_config)
+            quantization_config_copy.tokenizer = quantization_config.tokenizer or model_id
+            quantizer.quantize(ov_config=OVConfig(quantization_config=quantization_config_copy))
 
         return causal_model
 

optimum/intel/openvino/quantization.py

@@ -201,7 +201,7 @@ def from_pretrained(cls, model: PreTrainedModel, **kwargs):
     def quantize(
         self,
         calibration_dataset: Optional[Union[datasets.Dataset, nncf.Dataset, Iterable]] = None,
-        save_directory: Union[str, Path] = None,
+        save_directory: Optional[Union[str, Path]] = None,
         ov_config: OVConfig = None,
         file_name: Optional[str] = None,
         batch_size: int = 1,
@@ -217,7 +217,7 @@ def quantize(
             calibration_dataset (`datasets.Dataset` or `nncf.Dataset` or `Iterable`, *optional*):
                 A collection of data samples to use for quantization calibration. Is optional for weight-only
                 quantization and is required for full quantization.
-            save_directory (`Union[str, Path]`):
+            save_directory (`Union[str, Path]`, *optional*):
                 The directory where the quantized model should be saved.
             ov_config (`OVConfig`, *optional*):
                 The configuration containing the parameters related to quantization. If not provided, 8-bit symmetric
@@ -265,10 +265,6 @@ def quantize(
                 "as an instance of `OVWeightQuantizationConfig` for weight-only compression or as an instance of `OVQuantizationConfig` for full model quantization."
             )
 
-        if save_directory is None:
-            # TODO : can be set to self.model.config.name_or_path for OVModels when not provided
-            raise ValueError("`save_directory` needs to be specified")
-
         if ov_config is None:
             ov_config = OVConfig()
         if not isinstance(ov_config, OVConfig):
@@ -321,21 +317,41 @@ def quantize(
     def _quantize_ovbasemodel(
         self,
         ov_config: OVConfig,
-        save_directory: Union[str, Path],
+        save_directory: Union[str, Path] = None,
         calibration_dataset: Optional[Union[datasets.Dataset, nncf.Dataset, Iterable]] = None,
         batch_size: int = 1,
         data_collator: Optional[DataCollator] = None,
         remove_unused_columns: bool = True,
         **kwargs,
     ):
-        save_directory = Path(save_directory)
-        save_directory.mkdir(parents=True, exist_ok=True)
+        if save_directory is not None:
+            save_directory = Path(save_directory)
+            save_directory.mkdir(parents=True, exist_ok=True)
 
         quantization_config = ov_config.quantization_config
         if isinstance(quantization_config, OVWeightQuantizationConfig):
+            if calibration_dataset is None and isinstance(quantization_config.dataset, str):
+                from optimum.intel import OVModelForCausalLM
+
+                if isinstance(self.model, OVModelForCausalLM):
+                    from optimum.gptq.data import get_dataset, prepare_dataset
+
+                    tokenizer = AutoTokenizer.from_pretrained(quantization_config.tokenizer)
+                    nsamples = quantization_config.num_samples if quantization_config.num_samples else 128
+                    calibration_dataset = get_dataset(
+                        quantization_config.dataset, tokenizer, seqlen=32, nsamples=nsamples
+                    )
+                    calibration_dataset = prepare_dataset(calibration_dataset)
+                    calibration_dataset = nncf.Dataset(calibration_dataset, lambda x: self.model.prepare_inputs(**x))
+                else:
+                    raise ValueError(
+                        f"Can't create weight compression calibration dataset from string for {type(self.model)}"
+                    )
+
             _weight_only_quantization(self.model.model, quantization_config, calibration_dataset)
-            self.model.save_pretrained(save_directory)
-            ov_config.save_pretrained(save_directory)
+            if save_directory is not None:
+                self.model.save_pretrained(save_directory)
+                ov_config.save_pretrained(save_directory)
             return
         if not isinstance(quantization_config, OVQuantizationConfig):
             raise ValueError(f"Unsupported type of quantization config: {type(quantization_config)}")
@@ -387,8 +403,9 @@ def _quantize_ovbasemodel(
             **kwargs,
         )
         self.model.model = quantized_model
-        self.model.save_pretrained(save_directory)
-        ov_config.save_pretrained(save_directory)
+        if save_directory is not None:
+            self.model.save_pretrained(save_directory)
+            ov_config.save_pretrained(save_directory)
 
     def _quantize_torchmodel(
         self,
@@ -401,6 +418,10 @@ def _quantize_torchmodel(
         remove_unused_columns: bool = True,
         **kwargs,
     ):
+        if save_directory is None:
+            # TODO : can be set to self.model.config.name_or_path for OVModels when not provided
+            raise ValueError("`save_directory` needs to be specified")
+
         self._set_task()
         save_directory = Path(save_directory)
         save_directory.mkdir(parents=True, exist_ok=True)
@@ -660,14 +681,6 @@ def _weight_only_quantization(
             dataset = calibration_dataset
         else:
             dataset = nncf.Dataset(calibration_dataset)
-    elif config.dataset is not None and isinstance(config.dataset, str):
-        tokenizer = AutoTokenizer.from_pretrained(config.tokenizer)
-
-        from optimum.gptq.data import get_dataset, prepare_dataset
-
-        nsamples = config.num_samples if config.num_samples else 128
-        dataset = get_dataset(config.dataset, tokenizer, seqlen=32, nsamples=nsamples)
-        dataset = prepare_dataset(dataset)
 
     sensitivity_metric = None
     if isinstance(config.sensitivity_metric, str):