Fix config saving

Author: echarlaix

optimum/intel/openvino/configuration.py

@@ -78,25 +78,26 @@
 }
 
 
+
 DEFAULT_4BIT_CONFIGS = {
-    "databricks/dolly-v2-3b": {"mode": nncf.CompressWeightsMode.INT4_ASYM, "group_size": 32, "ratio": 0.5},
-    "EleutherAI/gpt-j-6b": {"mode": nncf.CompressWeightsMode.INT4_ASYM, "group_size": 64},
-    "facebook/opt-6.7b": {"mode": nncf.CompressWeightsMode.INT4_ASYM, "group_size": 64, "ratio": 0.8},
-    "bigscience/bloomz-7b1": {"mode": nncf.CompressWeightsMode.INT4_ASYM, "group_size": 32, "ratio": 0.6},
-    "togethercomputer/RedPajama-INCITE-7B-Instruct": {"mode": nncf.CompressWeightsMode.INT4_ASYM, "group_size": 128},
-    "HuggingFaceH4/zephyr-7b-beta": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 64, "ratio": 0.6},
-    "meta-llama/Llama-2-7b": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 128, "ratio": 0.6},
-    "meta-llama/Llama-2-7b-chat": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 128, "ratio": 0.8},
-    "meta-llama/Llama-2-13b-chat": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 64, "ratio": 0.8},
-    "stabilityai/stablelm-3b-4e1t": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 64, "ratio": 0.8},
-    "stablelm-epoch-3b-preview": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 64, "ratio": 0.8},
-    "stable-zephyr-3b-dpo": {"mode": nncf.CompressWeightsMode.INT4_ASYM, "group_size": 64, "ratio": 0.8},
-    "pansophic/rocket-3B": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 128, "ratio": 0.8},
-    "THUDM/chatglm2-6b": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 128, "ratio": 0.72},
-    "Qwen/Qwen-7B-Chat": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 128, "ratio": 0.6},
-    "openlm-research/open_llama_3b": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 64, "all_layers": True},
-    "tiiuae/falcon-7b": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 64, "all_layers": True},
-    "psmathur/orca_mini_3b": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 64, "all_layers": True},
+    "databricks/dolly-v2-3b": {"bits": 4, "sym": False, "group_size": 32, "ratio": 0.5},
+    "EleutherAI/gpt-j-6b": {"bits": 4, "sym": False, "group_size": 64},
+    "facebook/opt-6.7b": {"bits": 4, "sym": False, "group_size": 64, "ratio": 0.8},
+    "bigscience/bloomz-7b1": {"bits": 4, "sym": False, "group_size": 32, "ratio": 0.6},
+    "togethercomputer/RedPajama-INCITE-7B-Instruct": {"bits": 4, "sym": False, "group_size": 128},
+    "HuggingFaceH4/zephyr-7b-beta": {"bits": 4, "sym": True, "group_size": 64, "ratio": 0.6},
+    "meta-llama/Llama-2-7b": {"bits": 4, "sym": True, "group_size": 128, "ratio": 0.6},
+    "meta-llama/Llama-2-7b-chat": {"bits": 4, "sym": True, "group_size": 128, "ratio": 0.8},
+    "meta-llama/Llama-2-13b-chat": {"bits": 4, "sym": True, "group_size": 64, "ratio": 0.8},
+    "stabilityai/stablelm-3b-4e1t": {"bits": 4, "sym": True, "group_size": 64, "ratio": 0.8},
+    "stablelm-epoch-3b-preview": {"bits": 4, "sym": True, "group_size": 64, "ratio": 0.8},
+    "stable-zephyr-3b-dpo": {"bits": 4, "sym": False, "group_size": 64, "ratio": 0.8},
+    "pansophic/rocket-3B": {"bits": 4, "sym": True, "group_size": 128, "ratio": 0.8},
+    "THUDM/chatglm2-6b": {"bits": 4, "sym": True, "group_size": 128, "ratio": 0.72},
+    "Qwen/Qwen-7B-Chat": {"bits": 4, "sym": True, "group_size": 128, "ratio": 0.6},
+    "openlm-research/open_llama_3b": {"bits": 4, "sym": True, "group_size": 64, "all_layers": True},
+    "tiiuae/falcon-7b": {"bits": 4, "sym": True, "group_size": 64, "all_layers": True},
+    "psmathur/orca_mini_3b": {"bits": 4, "sym": True, "group_size": 64, "all_layers": True},
 }
 
 
@@ -159,8 +160,11 @@ class OVWeightQuantizationConfig(QuantizationConfigMixin):
     loaded using `optimum-intel` api for quantization with NNCF.
 
     Args:
-        mode (`nncf.CompressWeightsMode`, *optional*, defaults to INT8_ASYM):
-            The model defines the weight compressoin method (4-bit, 8-bit, etc.) available in nncf.compress_weights nncf.CompressWeightsMode.
+
+        bits (`int`, defaults to 8):
+            The number of bits to quantize to.
+        sym (`bool`, *optional*, defaults to `False`):
+            Whether to use symetric quantization.
         tokenizer (`str` or `PreTrainedTokenizerBase`, *optional*):
             The tokenizer used to process the dataset. You can pass either:
                 - A custom tokenizer object.
@@ -191,26 +195,27 @@ class OVWeightQuantizationConfig(QuantizationConfigMixin):
 
     def __init__(
         self,
-        mode=None,
+        bits: int = 8,
+        sym: bool = False,
         tokenizer: Any = None,
-        dataset: Optional[Union[nncf.Dataset, str]] = None,
+        dataset: Optional[str] = None,
         ratio: Optional[float] = None,
         group_size: Optional[int] = None,
         all_layers: Optional[bool] = None,
-        sensitivity_metric: Optional[nncf.SensitivityMetric] = None,
-        awq: Optional[bool] = None,
-        ignored_scope: Optional[nncf.IgnoredScope] = None,
+        sensitivity_metric: Optional[str] = None,
+        ignored_scope: Optional[dict] = None,
         **kwargs,
     ):
-        self.mode = mode
+        self.bits = bits
+        self.sym = sym
         self.tokenizer = tokenizer
         self.dataset = dataset
         self.group_size = group_size
         self.ratio = ratio
-        self.ignored_scope = ignored_scope
         self.all_layers = all_layers
         self.sensitivity_metric = sensitivity_metric
-        self.awq = awq
+        self.ignored_scope = ignored_scope
+        self.quant_method = "default" # TODO : enable AWQ after nncf v2.9.0 release
         self.post_init()
 
     def post_init(self):
@@ -229,5 +234,9 @@ def post_init(self):
                 )
 
 
+        if self.bits not in [4, 8]:
+            raise ValueError(f"Only support quantization to [4,8] bits but found {self.bits}")
+
+
 def _check_default_4bit_configs(config: PretrainedConfig):
     return DEFAULT_4BIT_CONFIGS.get(config.name_or_path, None)

optimum/intel/openvino/quantization.py

@@ -24,10 +24,11 @@
 import transformers
 from accelerate.data_loader import DataLoaderStateMixin
 from datasets import Dataset, load_dataset
-from nncf import NNCFConfig
+from nncf import NNCFConfig, CompressWeightsMode, SensitivityMetric, IgnoredScope
 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
+
 from openvino._offline_transformations import compress_quantize_weights_transformation
 from openvino.runtime import Core, Tensor
 from torch.utils._pytree import tree_map
@@ -54,7 +55,9 @@
 )
 
 
-COMPRESSION_OPTIONS = {
+
+# TODO : remove as unused
+_COMPRESSION_OPTIONS = {
     "int8": {"mode": nncf.CompressWeightsMode.INT8},
     "int4_sym_g128": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 128},
     "int4_asym_g128": {"mode": nncf.CompressWeightsMode.INT4_ASYM, "group_size": 128},
@@ -272,12 +275,11 @@ def quantize(
     def _get_compression_options(self, config: OVConfig):
         options = {}
         if config is not None and "type" in config.compression:
-            options = COMPRESSION_OPTIONS[config.compression["type"]]
+            options = _COMPRESSION_OPTIONS[config.compression["type"]]
             if "ratio" in config.compression:
                 options["ratio"] = config.compression["ratio"]
         return options
 
-    # TODO : add ov_config
     def _quantize_ovbasemodel(
         self,
         calibration_dataset: Dataset,
@@ -597,6 +599,7 @@ def _int4_weight_only_quantization(
             config = OVWeightQuantizationConfig.from_dict(quantization_config)
 
         dataset = config.dataset
+
         if config.dataset is not None and isinstance(config.dataset, str):
             tokenizer = config.tokenizer
             if tokenizer is None:
@@ -610,15 +613,24 @@ def _int4_weight_only_quantization(
             dataset = prepare_dataset(dataset)
             dataset = nncf.Dataset(dataset, lambda x: model.prepare_inputs(**x))
 
+
+        sensitivity_metric = None
+        if isinstance(config.sensitivity_metric, str):
+            sensitivity_metric = getattr(SensitivityMetric, config.sensitivity_metric.upper())
+
+        ignored_scope = None
+        if isinstance(config.ignored_scope, dict):
+            ignored_scope = IgnoredScope(**config.ignored_scope)
+
         model.model = nncf.compress_weights(
             ov_model,
-            mode=config.mode,
+            mode=CompressWeightsMode.INT4_SYM if config.sym else CompressWeightsMode.INT4_ASYM,
             ratio=config.ratio,
             group_size=config.group_size,
             all_layers=config.all_layers,
-            sensitivity_metric=config.sensitivity_metric,
-            awq=config.awq,
-            ignored_scope=config.ignored_scope,
+            sensitivity_metric=sensitivity_metric,
+            # awq=config.quant_method == "awq", # TODO : remove and add it back once nncf v2.9.0
+            ignored_scope=ignored_scope,
             dataset=dataset,
         )
     else: