[OpenVINO] Add workaround logic for default int4 quantization of openai/gpt-oss-20b model (#1490)

* Add a workaround for GPT-OSS quantization

* Fix tests

* Add exception

* Address review comments

* Add config save/load tests

Author: nikita-savelyevv

optimum/exporters/openvino/__main__.py

@@ -520,9 +520,18 @@ class StoreAttr(object):
                     "Quantization of the weights requires nncf, please install it with `pip install nncf`"
                 )
 
+            from optimum.intel.openvino.configuration import _GPTOSSQuantizationConfig
             from optimum.intel.openvino.quantization import _weight_only_quantization
 
-            _weight_only_quantization(submodel, quantization_config)
+            if isinstance(quantization_config, _GPTOSSQuantizationConfig):
+                # A workaround for GPT-OSS model is required to run quantization twice, this way it is possible to
+                # selectively quantize some weights to 4 bits and some to 8 bits.
+                _weight_only_quantization(submodel, quantization_config.quantization_config1)
+                _weight_only_quantization(
+                    submodel, quantization_config.quantization_config2, verify_not_optimized=False
+                )
+            else:
+                _weight_only_quantization(submodel, quantization_config)
             compressed_submodel_path = submodel_path.parent / f"{submodel_path.stem}_compressed.xml"
             save_model(submodel, compressed_submodel_path, compress_to_fp16=False)
             del submodel

optimum/intel/openvino/configuration.py

@@ -316,13 +316,34 @@ class OVQuantizationMethod(str, Enum):
         },
     },
     "openai/gpt-oss-20b": {
-        "bits": 4,
-        "sym": True,
-        "group_size": 32,
-        "ignored_scope": {
-            "patterns": [".*self_attn.*", ".*router.*"],
+        "quantization_config1": {
+            "bits": 4,
+            "sym": True,
+            "group_size": 32,
+            # With ignored scope below we keep some weights in their original precision during the first quantization
+            # run and then quantize them to int8 in the second run.
+            "ignored_scope": {"patterns": [".*self_attn.*", ".*router.*"]},
+        },
+        "quantization_config2": {
+            "bits": 8,
+            "sym": False,
+            "weight_only": True,
+        },
+    },
+    "openai/gpt-oss-120b": {
+        "quantization_config1": {
+            "bits": 4,
+            "sym": True,
+            "group_size": 32,
+            # With ignored scope below we keep some weights in their original precision during the first quantization
+            # run and then quantize them to int8 in the second run.
+            "ignored_scope": {"patterns": [".*self_attn.*", ".*router.*"]},
+        },
+        "quantization_config2": {
+            "bits": 8,
+            "sym": False,
+            "weight_only": True,
         },
-        "backup_precision": "none",
     },
 }
 
@@ -1149,6 +1170,8 @@ def _get_dtype(quantization_config):
         elif isinstance(quantization_config, OVPipelineQuantizationConfig):
             dtypes = [OVConfig._get_dtype(config) for config in quantization_config.quantization_configs.values()]
             dtype = "_".join(dtypes)
+        elif isinstance(quantization_config, _GPTOSSQuantizationConfig):
+            dtype = "int4_int8"
         else:
             raise ValueError(f"Unsupported type of quantization config: {type(quantization_config)}")
         return dtype
@@ -1366,6 +1389,36 @@ def post_init(self):
             submodel_config.post_init()
 
 
+class _GPTOSSQuantizationConfig(QuantizationConfigMixin):
+    def __init__(
+        self,
+        quantization_config1: Union[Dict, OVWeightQuantizationConfig],
+        quantization_config2: Union[Dict, OVWeightQuantizationConfig],
+        **kwargs,
+    ):
+        """
+        Configuration class for GPT-OSS quantization.
+
+        # TODO (nikita.savelyevv): Introduce OVSequentialQuantizationConfig to support this.
+        """
+
+        if isinstance(quantization_config1, dict):
+            quantization_config1 = OVWeightQuantizationConfig.from_dict(quantization_config1)
+        self.quantization_config1 = quantization_config1
+        self.quantization_config1.post_init()
+
+        if isinstance(quantization_config2, dict):
+            quantization_config2 = OVWeightQuantizationConfig.from_dict(quantization_config2)
+        self.quantization_config2 = quantization_config2
+        self.quantization_config2.post_init()
+
+    def to_dict(self) -> Dict[str, Any]:
+        result = super().to_dict()
+        result["quantization_config1"] = self.quantization_config1.to_dict()
+        result["quantization_config2"] = self.quantization_config2.to_dict()
+        return result
+
+
 def _quantization_config_from_dict(config_dict: Dict[str, Any]) -> OVQuantizationConfigBase:
     """
     Helper function to create a quantization config from a dictionary.
@@ -1379,6 +1432,10 @@ def _quantization_config_from_dict(config_dict: Dict[str, Any]) -> OVQuantizatio
     if "quantization_configs" in config_dict:
         return OVPipelineQuantizationConfig.from_dict(config_dict)
 
+    # Check for GPT-OSS quantization config
+    if "quantization_config1" in config_dict and "quantization_config2" in config_dict:
+        return _GPTOSSQuantizationConfig.from_dict(config_dict)
+
     # Either OVWeightQuantizationConfig or OVQuantizationConfig
     # Try to detect the type of config based on the keys present in the dictionary
     wq_args = set(inspect.getfullargspec(OVWeightQuantizationConfig.__init__).args)

optimum/intel/openvino/modeling_decoder.py

@@ -876,7 +876,14 @@ def _from_pretrained(
             init_cls = cls
 
         if isinstance(quantization_config, dict) and quantization_config == {"bits": 4}:
-            default_config = get_default_quantization_config(config.name_or_path, weight_format="int4")
+            if config.name_or_path in ["openai/gpt-oss-20b", "openai/gpt-oss-120b"]:
+                raise NotImplementedError(
+                    "Quantization with the default 4-bit config is not supported through Python API for openai/gpt-oss-20b model. "
+                    "Please export the model via optimum-cli with `--weight-format int4` argument. This way the "
+                    "recommended quantization config will be used."
+                )
+            else:
+                default_config = get_default_quantization_config(config.name_or_path, weight_format="int4")
             quantization_config = cls._prepare_quantization_config(
                 default_config or _DEFAULT_4BIT_WQ_CONFIG, load_in_8bit
             )

optimum/intel/openvino/quantization.py

@@ -1561,9 +1561,11 @@ def _weight_only_quantization(
     model: openvino.Model,
     quantization_config: Union[OVWeightQuantizationConfig, Dict],
     calibration_dataset: Optional[Union[nncf.Dataset, Iterable]] = None,
+    verify_not_optimized: bool = True,
     **kwargs,
 ) -> openvino.Model:
-    _verify_not_optimized(model)
+    if verify_not_optimized:
+        _verify_not_optimized(model)
     config = quantization_config
     if isinstance(config, dict):
         config = OVWeightQuantizationConfig.from_dict(quantization_config)

tests/openvino/test_exporters_cli.py

@@ -1108,26 +1108,47 @@ def test_exporters_cli_full_quantization(
                 expected_fake_nodes_per_model,
             )
 
-    @parameterized.expand(
-        [
-            (
-                "falcon-40b",
-                "bigscience/bloomz-560m",
-                AutoModelForCausalLM,
-                OVModelForCausalLM,
-                "--task text-generation-with-past --weight-format int4",
-                _DEFAULT_4BIT_WQ_CONFIGS,
-            ),
-            (
-                "clip",
-                "hf-tiny-model-private/tiny-random-CLIPModel",
-                AutoModelForZeroShotImageClassification,
-                OVModelForZeroShotImageClassification,
-                "--task zero-shot-image-classification --quant-mode int8",
-                _DEFAULT_INT8_FQ_CONFIGS,
-            ),
-        ]
-    )
+    DEFAULT_CONFIG_TEST_CONFIGURATIONS = [
+        (
+            "falcon-40b",
+            "bigscience/bloomz-560m",
+            AutoModelForCausalLM,
+            OVModelForCausalLM,
+            "--task text-generation-with-past --weight-format int4",
+            _DEFAULT_4BIT_WQ_CONFIGS,
+            {"model": {"int8": 6, "int4": 6}},
+            {"model": 0},
+        ),
+        (
+            "clip",
+            "hf-tiny-model-private/tiny-random-CLIPModel",
+            AutoModelForZeroShotImageClassification,
+            OVModelForZeroShotImageClassification,
+            "--task zero-shot-image-classification --quant-mode int8",
+            _DEFAULT_INT8_FQ_CONFIGS,
+            {"model": {"int8": 65}},
+            {"model": 65},
+        ),
+        (
+            "gpt_oss_mxfp4",
+            "openai/gpt-oss-20b",
+            AutoModelForCausalLM,
+            OVModelForCausalLM,
+            "--task text-generation-with-past --weight-format int4",
+            _DEFAULT_4BIT_WQ_CONFIGS,
+            {"model": {"int8": 22, "int4": 4}},
+            {"model": 0},
+        ),
+    ]
+
+    # filter models type depending on min max transformers version
+    SUPPORTED_DEFAULT_CONFIG_TEST_CONFIGURATIONS = [
+        config
+        for config in DEFAULT_CONFIG_TEST_CONFIGURATIONS
+        if TEST_NAME_TO_MODEL_TYPE.get(config[0], config[0]) in get_supported_model_for_library("transformers")
+    ]
+
+    @parameterized.expand(SUPPORTED_DEFAULT_CONFIG_TEST_CONFIGURATIONS)
     def test_exporters_cli_with_default_config(
         self,
         model_name,
@@ -1136,6 +1157,8 @@ def test_exporters_cli_with_default_config(
         ov_model_cls,
         options,
         default_configs_collection,
+        expected_num_weight_nodes_per_model,
+        expected_fake_nodes_per_model,
     ):
         with TemporaryDirectory() as tmpdir:
             pt_model = auto_model_cls.from_pretrained(MODEL_NAMES[model_name])
@@ -1167,15 +1190,26 @@ def test_exporters_cli_with_default_config(
             )
 
             model = ov_model_cls.from_pretrained(tmpdir)
+
+            check_compression_state_per_model(
+                self,
+                model.ov_submodels,
+                expected_num_weight_nodes_per_model,
+                expected_fake_nodes_per_model,
+            )
+
             rt_info = model.model.get_rt_info()
             nncf_info = rt_info["nncf"]
             model_quantization_config = nncf_info["weight_compression" if is_weight_compression else "quantization"]
 
             default_config = {**default_configs_collection[model_id]}
-            default_config.pop("dataset", None)
+            if "quantization_config2" in default_config:
+                # For GPT-OSS use the second config as reference
+                default_config = default_config["quantization_config2"]
+            dataset = default_config.pop("dataset", None)
+            default_config.pop("weight_only", None)
             if is_weight_compression:
                 bits = default_config.pop("bits", None)
-                self.assertEqual(bits, 4)
                 sym = default_config.pop("sym", False)
                 default_config["mode"] = f"int{bits}_{'sym' if sym else 'asym'}"
                 quant_method = default_config.pop("quant_method", None)
@@ -1184,7 +1218,8 @@ def test_exporters_cli_with_default_config(
                 advanced_parameters = eval(model_quantization_config["advanced_parameters"].value)
                 model_quantization_config["statistics_path"] = Mock()
                 model_quantization_config["statistics_path"].value = advanced_parameters["statistics_path"]
-                default_config["statistics_path"] = f"{tmpdir}/statistics"
+                if dataset is not None:
+                    default_config["statistics_path"] = f"{tmpdir}/statistics"
             else:
                 dtype = default_config.pop("dtype", None)
                 self.assertEqual(dtype, "int8")

tests/openvino/test_quantization.py

@@ -78,6 +78,7 @@
     _DEFAULT_4BIT_WQ_CONFIGS,
     _DEFAULT_4BIT_WQ_CONFIG,
     _quantization_config_from_dict,
+    _GPTOSSQuantizationConfig,
 )
 from optimum.intel.openvino.modeling_visual_language import _OVNanoLlavaForCausalLM
 from optimum.intel.openvino.utils import TemporaryDirectory
@@ -1899,6 +1900,12 @@ class OVQuantizationConfigTest(unittest.TestCase):
                 advanced_parameters=nncf.AdvancedCompressionParameters(),
             ),
         ),
+        (
+            _GPTOSSQuantizationConfig(
+                quantization_config1=OVWeightQuantizationConfig(bits=4, group_size=16),
+                quantization_config2=OVWeightQuantizationConfig(bits=8),
+            ),
+        ),
     )
 
     QUANTIZATION_CONFIG_DICTS = (
@@ -1988,6 +1995,14 @@ class OVQuantizationConfigTest(unittest.TestCase):
             OVPipelineQuantizationConfig,
             None,
         ),
+        (
+            dict(
+                quantization_config1=dict(bits=4, group_size=16),
+                quantization_config2=dict(bits=8, weight_only=True),
+            ),
+            _GPTOSSQuantizationConfig,
+            None,
+        ),
     )
 
     QUANTIZATION_CONFIGS_WITH_KWARGS = (