[OV] Add support for nf4_f8e4m3 quantization mode (#1148)

* Initial commit

* Fix tests

* Add test

* OVMixedQuantizationConfig proposal

* Polishing changes

* Style

* Fix

* Hybrid quantization as mixed quantization

* Fix

* Fix

* Renaming + docstrings

* Update num_samples default value

* Removed ignored scope base class; fix how mixed precision quantization is applied

* Rename weight & activation format to dtype

* Add int4_f8e4m3 quant mode

* Update description

* Add 'nf4_f8e5m2', 'int4_f8e5m2'; add backup precision

* Address comments

* Trigger Test

* Address comments

---------

Co-authored-by: Nikita Malinin <[email protected]>

Author: nikita-savelyevv

docs/source/openvino/export.mdx

@@ -31,7 +31,8 @@ Check out the help for more options:
 
 ```text
 usage: optimum-cli export openvino [-h] -m MODEL [--task TASK] [--framework {pt,tf}] [--trust-remote-code]
-                                   [--weight-format {fp32,fp16,int8,int4,mxfp4,nf4}] [--quant-mode {int8,f8e4m3,f8e5m2}]
+                                   [--weight-format {fp32,fp16,int8,int4,mxfp4,nf4}]
+                                   [--quant-mode {int8,f8e4m3,f8e5m2,nf4_f8e4m3,nf4_f8e5m2,int4_f8e4m3,int4_f8e5m2}]
                                    [--library {transformers,diffusers,timm,sentence_transformers,open_clip}]
                                    [--cache_dir CACHE_DIR] [--pad-token-id PAD_TOKEN_ID] [--ratio RATIO] [--sym]
                                    [--group-size GROUP_SIZE] [--backup-precision {none,int8_sym,int8_asym}]
@@ -67,7 +68,7 @@ Optional arguments:
                         on your local machine arbitrary code present in the model repository.
   --weight-format {fp32,fp16,int8,int4,mxfp4,nf4}
                         The weight format of the exported model.
-  --quant-mode {int8,f8e4m3,f8e5m2}
+  --quant-mode {int8,f8e4m3,f8e5m2,nf4_f8e4m3,nf4_f8e5m2,int4_f8e4m3,int4_f8e5m2}
                         Quantization precision mode. This is used for applying full model quantization including
                         activations.
   --library {transformers,diffusers,timm,sentence_transformers,open_clip}

optimum/commands/export/openvino.py

@@ -78,7 +78,7 @@ def parse_args_openvino(parser: "ArgumentParser"):
     optional_group.add_argument(
         "--quant-mode",
         type=str,
-        choices=["int8", "f8e4m3", "f8e5m2"],
+        choices=["int8", "f8e4m3", "f8e5m2", "nf4_f8e4m3", "nf4_f8e5m2", "int4_f8e4m3", "int4_f8e5m2"],
         default=None,
         help=(
             "Quantization precision mode. This is used for applying full model quantization including activations. "
@@ -352,23 +352,7 @@ def run(self):
                 if no_compression_parameter_provided(self.args) and self.args.weight_format == "int4":
                     quantization_config = get_default_int4_config(self.args.model)
                 else:
-                    is_int8 = self.args.weight_format == "int8"
-                    quantization_config = {
-                        "bits": 8 if is_int8 else 4,
-                        "ratio": 1.0 if is_int8 else (self.args.ratio or _DEFAULT_4BIT_CONFIG["ratio"]),
-                        "sym": self.args.sym or False,
-                        "group_size": -1 if is_int8 else self.args.group_size,
-                        "all_layers": None if is_int8 else self.args.all_layers,
-                        "dataset": self.args.dataset,
-                        "num_samples": self.args.num_samples,
-                        "quant_method": "awq" if self.args.awq else "default",
-                        "sensitivity_metric": self.args.sensitivity_metric,
-                        "scale_estimation": self.args.scale_estimation,
-                        "gptq": self.args.gptq,
-                        "lora_correction": self.args.lora_correction,
-                        "weight_format": self.args.weight_format,
-                        "backup_precision": self.args.backup_precision,
-                    }
+                    quantization_config = prepare_wc_config(self.args, _DEFAULT_4BIT_CONFIG)
 
                 if quantization_config.get("dataset", None) is not None:
                     quantization_config["trust_remote_code"] = self.args.trust_remote_code
@@ -378,16 +362,24 @@ def run(self):
                     raise ValueError(
                         "Dataset is required for full quantization. Please provide it with --dataset argument."
                     )
-                quantization_config = {
-                    "weight_format": self.args.quant_mode,
-                    "activation_format": self.args.quant_mode,
-                    "bits": 8,
-                    "sym": self.args.sym or False,
-                    "dataset": self.args.dataset,
-                    "num_samples": self.args.num_samples,
-                    "smooth_quant_alpha": self.args.smooth_quant_alpha,
-                    "trust_remote_code": self.args.trust_remote_code,
-                }
+
+                if self.args.quant_mode in ["nf4_f8e4m3", "nf4_f8e5m2", "int4_f8e4m3", "int4_f8e5m2"]:
+                    wc_config = prepare_wc_config(self.args, _DEFAULT_4BIT_CONFIG)
+                    wc_dtype, q_dtype = self.args.quant_mode.split("_")
+                    wc_config["dtype"] = wc_dtype
+
+                    q_config = prepare_q_config(self.args)
+                    q_config["dtype"] = q_dtype
+
+                    quantization_config = {
+                        "weight_quantization_config": wc_config,
+                        "full_quantization_config": q_config,
+                        "num_samples": self.args.num_samples,
+                        "dataset": self.args.dataset,
+                        "trust_remote_code": self.args.trust_remote_code,
+                    }
+                else:
+                    quantization_config = prepare_q_config(self.args)
                 ov_config = OVConfig(quantization_config=quantization_config)
 
         quantization_config = ov_config.quantization_config if ov_config else None
@@ -486,3 +478,35 @@ def run(self):
                 variant=self.args.variant,
                 # **input_shapes,
             )
+
+
+def prepare_wc_config(args, default_configs):
+    is_int8 = args.weight_format == "int8"
+    return {
+        "bits": 8 if is_int8 else 4,
+        "ratio": 1.0 if is_int8 else (args.ratio or default_configs["ratio"]),
+        "sym": args.sym or False,
+        "group_size": -1 if is_int8 else args.group_size,
+        "all_layers": None if is_int8 else args.all_layers,
+        "dataset": args.dataset,
+        "num_samples": args.num_samples,
+        "quant_method": "awq" if args.awq else "default",
+        "sensitivity_metric": args.sensitivity_metric,
+        "scale_estimation": args.scale_estimation,
+        "gptq": args.gptq,
+        "lora_correction": args.lora_correction,
+        "dtype": args.weight_format,
+        "backup_precision": args.backup_precision,
+    }
+
+
+def prepare_q_config(args):
+    return {
+        "dtype": args.quant_mode,
+        "bits": 8,
+        "sym": args.sym or False,
+        "dataset": args.dataset,
+        "num_samples": args.num_samples,
+        "smooth_quant_alpha": args.smooth_quant_alpha,
+        "trust_remote_code": args.trust_remote_code,
+    }

optimum/intel/__init__.py

@@ -84,6 +84,7 @@
             "OVQuantizationConfig",
             "OVWeightQuantizationConfig",
             "OVDynamicQuantizationConfig",
+            "OVMixedQuantizationConfig",
         ]
     )
 else:
@@ -94,6 +95,7 @@
             "OVQuantizationConfig",
             "OVWeightQuantizationConfig",
             "OVDynamicQuantizationConfig",
+            "OVMixedQuantizationConfig",
         ]
     )
 
@@ -272,6 +274,7 @@
     except OptionalDependencyNotAvailable:
         from .utils.dummy_openvino_and_nncf_objects import (
             OVDynamicQuantizationConfig,
+            OVMixedQuantizationConfig,
             OVQuantizationConfig,
             OVQuantizer,
             OVTrainingArguments,
@@ -280,6 +283,7 @@
     else:
         from .openvino import (
             OVDynamicQuantizationConfig,
+            OVMixedQuantizationConfig,
             OVQuantizationConfig,
             OVQuantizer,
             OVTrainingArguments,

optimum/intel/openvino/__init__.py

@@ -55,7 +55,13 @@
         from .trainer import OVTrainer
 
 
-from .configuration import OVConfig, OVDynamicQuantizationConfig, OVQuantizationConfig, OVWeightQuantizationConfig
+from .configuration import (
+    OVConfig,
+    OVDynamicQuantizationConfig,
+    OVMixedQuantizationConfig,
+    OVQuantizationConfig,
+    OVWeightQuantizationConfig,
+)
 from .modeling import (
     OVModelForAudioClassification,
     OVModelForAudioFrameClassification,