Add white_list to config (#1430)

* add white list
---------

Signed-off-by: yiliu30 <[email protected]>
Signed-off-by: chensuyue <[email protected]>
Co-authored-by: chensuyue <[email protected]>

Author: yiliu30

.azure-pipelines/model-test.yml

@@ -11,13 +11,13 @@ pr:
       - neural_compressor
       - setup.py
       - requirements.txt
-      - .azure-pipelines/model-test.yml
       - .azure-pipelines/scripts/models
       - examples/tensorflow/oob_models/quantization/ptq
     exclude:
       - test
       - neural_compressor/common
       - neural_compressor/torch
+      - neural_compressor/tensorflow
 
 pool: MODEL_PERF_TEST_TF
 

.azure-pipelines/ut-basic-no-cover.yml

@@ -17,6 +17,7 @@ pr:
       - test/3x
       - neural_compressor/common
       - neural_compressor/torch
+      - neural_compressor/tensorflow
 
 pool: ICX-16C
 

.azure-pipelines/ut-basic.yml

@@ -17,6 +17,7 @@ pr:
       - test/3x
       - neural_compressor/common
       - neural_compressor/torch
+      - neural_compressor/tensorflow
 
 pool: ICX-16C
 

.azure-pipelines/ut-itrex.yml

@@ -12,7 +12,10 @@ pr:
       - setup.py
       - requirements.txt
       - .azure-pipelines/scripts/ut/run_itrex.sh
-      - .azure-pipelines/ut-itrex.yml
+    exclude:
+      - neural_compressor/common
+      - neural_compressor/torch
+      - neural_compressor/tensorflow
 
 pool: MODEL_PERF_TEST
 

neural_compressor/common/base_config.py

@@ -23,7 +23,15 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from neural_compressor.common.logger import Logger
-from neural_compressor.common.utility import BASE_CONFIG, COMPOSABLE_CONFIG, GLOBAL, LOCAL
+from neural_compressor.common.utility import (
+    BASE_CONFIG,
+    COMPOSABLE_CONFIG,
+    DEFAULT_WHITE_LIST,
+    EMPTY_WHITE_LIST,
+    GLOBAL,
+    LOCAL,
+    OP_NAME_OR_MODULE_TYPE,
+)
 
 logger = Logger().get_logger()
 
@@ -59,18 +67,43 @@ class BaseConfig(ABC):
     """The base config for all algorithm configs."""
 
     name = BASE_CONFIG
+    params_list = []
 
-    def __init__(self) -> None:
+    def __init__(self, white_list: Optional[List[OP_NAME_OR_MODULE_TYPE]] = DEFAULT_WHITE_LIST) -> None:
         self._global_config: Optional[BaseConfig] = None
         # For PyTorch, operator_type is the collective name for module type and functional operation type,
         # for example, `torch.nn.Linear`, and `torch.nn.functional.linear`.
         # local config is the collections of operator_type configs and operator configs
         self._local_config: Dict[str, Optional[BaseConfig]] = {}
+        self._white_list = white_list
+
+    def _post_init(self):
+        if self.white_list == DEFAULT_WHITE_LIST:
+            global_config = self.get_params_dict()
+            self._global_config = self.__class__(**global_config, white_list=None)
+        elif isinstance(self.white_list, list) and len(self.white_list) > 0:
+            for op_name_or_type in self.white_list:
+                global_config = self.get_params_dict()
+                tmp_config = self.__class__(**global_config, white_list=None)
+                self.set_local(op_name_or_type, tmp_config)
+        elif self.white_list == EMPTY_WHITE_LIST:
+            return
+        else:
+            raise NotImplementedError(
+                f"The white list should be one of {DEFAULT_WHITE_LIST}, {EMPTY_WHITE_LIST},"
+                " a not empty list, but got {self.white_list}"
+            )
+
+    @property
+    def white_list(self):
+        return self._white_list
+
+    @white_list.setter
+    def white_list(self, op_name_or_type_list: Optional[List[OP_NAME_OR_MODULE_TYPE]]):
+        self._white_list = op_name_or_type_list
 
     @property
     def global_config(self):
-        if self._global_config is None:
-            self._global_config = self.__class__(**self.to_dict())
         return self._global_config
 
     @global_config.setter
@@ -88,25 +121,28 @@ def local_config(self, config):
     def set_local(self, operator_name: str, config: BaseConfig) -> BaseConfig:
         if operator_name in self.local_config:
             logger.warning("The configuration for %s has already been set, update it.", operator_name)
-        if self.global_config is None:
-            self.global_config = self.__class__(**self.to_dict())
         self.local_config[operator_name] = config
         return self
 
     def to_dict(self, params_list=[], operator2str=None):
         result = {}
-        global_config = {}
-        for param in params_list:
-            global_config[param] = getattr(self, param)
+        global_config = self.get_params_dict()
         if bool(self.local_config):
             result[LOCAL] = {}
             for op_name, config in self.local_config.items():
                 result[LOCAL][op_name] = config.to_dict()
-            result[GLOBAL] = global_config
+            if self.global_config:
+                result[GLOBAL] = global_config
         else:
             result = global_config
         return result
 
+    def get_params_dict(self):
+        result = dict()
+        for param in self.params_list:
+            result[param] = getattr(self, param)
+        return result
+
     @classmethod
     def from_dict(cls, config_dict, str2operator=None):
         """Construct config from a dict.
@@ -205,7 +241,8 @@ def to_config_mapping(
             global_config = config.global_config
             op_type_config_dict, op_name_config_dict = config._get_op_name_op_type_config()
             for op_name, op_type in model_info:
-                config_mapping[(op_type, op_name)] = global_config
+                if self.global_config is not None:
+                    config_mapping[(op_type, op_name)] = global_config
                 if op_type in op_type_config_dict:
                     config_mapping[(op_type, op_name)] = op_name_config_dict[op_type]
                 if op_name in op_name_config_dict:

neural_compressor/common/utility.py

@@ -21,11 +21,17 @@
 # constants for configs
 GLOBAL = "global"
 LOCAL = "local"
+DEFAULT_WHITE_LIST = "*"
+EMPTY_WHITE_LIST = None
 
 # config name
 BASE_CONFIG = "base_config"
 COMPOSABLE_CONFIG = "composable_config"
 RTN_WEIGHT_ONLY_QUANT = "rtn_weight_only_quant"
 STATIC_QUANT = "static_quant"
 GPTQ = "gptq"
-DUMMY_CONFIG = "dummy_config"
+
+
+from typing import Callable, Union
+
+OP_NAME_OR_MODULE_TYPE = Union[str, Callable]

neural_compressor/tensorflow/quantization/config.py

@@ -18,12 +18,12 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Callable, Dict, List, NamedTuple, Union
+from typing import Callable, Dict, List, NamedTuple, Optional, Union
 
 import tensorflow as tf
 
 from neural_compressor.common.base_config import BaseConfig, register_config, registered_configs
-from neural_compressor.common.utility import STATIC_QUANT
+from neural_compressor.common.utility import DEFAULT_WHITE_LIST, OP_NAME_OR_MODULE_TYPE, STATIC_QUANT
 
 FRAMEWORK_NAME = "keras"
 
@@ -89,6 +89,7 @@ def __init__(
         act_dtype: str = "int8",
         act_sym: bool = True,
         act_granularity: str = "per_tensor",
+        white_list: Optional[List[OP_NAME_OR_MODULE_TYPE]] = DEFAULT_WHITE_LIST,
     ):
         """Init static quantization config.
 
@@ -100,13 +101,14 @@ def __init__(
             act_sym (bool): Indicates whether activations are symmetric, default is True.
             act_granularity (str): Calculate tensor-wise scales or channel-wise scales for activations.
         """
-        super().__init__()
+        super().__init__(white_list=white_list)
         self.weight_dtype = weight_dtype
         self.weight_sym = weight_sym
         self.weight_granularity = weight_granularity
         self.act_dtype = act_dtype
         self.act_sym = act_sym
         self.act_granularity = act_granularity
+        self._post_init()
 
     def to_dict(self):
         return super().to_dict(params_list=self.params_list, operator2str=operator2str)

neural_compressor/torch/__init__.py

@@ -19,8 +19,6 @@
     quantize,
     RTNWeightQuantConfig,
     get_default_rtn_config,
-    DummyConfig,
-    get_default_dummy_config,
     GPTQConfig,
     get_default_gptq_config,
 )

neural_compressor/torch/quantization/__init__.py

@@ -16,8 +16,6 @@
 from neural_compressor.torch.quantization.config import (
     RTNWeightQuantConfig,
     get_default_rtn_config,
-    DummyConfig,
-    get_default_dummy_config,
     GPTQConfig,
     get_default_gptq_config,
 )

neural_compressor/torch/quantization/config.py

@@ -18,12 +18,12 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import Callable, Dict, List, NamedTuple, Union
+from typing import Callable, Dict, List, NamedTuple, Optional, Union
 
 import torch
 
 from neural_compressor.common.base_config import BaseConfig, register_config, registered_configs
-from neural_compressor.common.utility import DUMMY_CONFIG, GPTQ, RTN_WEIGHT_ONLY_QUANT
+from neural_compressor.common.utility import DEFAULT_WHITE_LIST, GPTQ, OP_NAME_OR_MODULE_TYPE, RTN_WEIGHT_ONLY_QUANT
 
 FRAMEWORK_NAME = "torch"
 
@@ -87,6 +87,7 @@ def __init__(
         double_quant_bits: int = 8,
         double_quant_sym: bool = True,
         double_quant_group_size: int = 256,
+        white_list: Optional[List[OP_NAME_OR_MODULE_TYPE]] = DEFAULT_WHITE_LIST,
     ):
         """Init RTN weight-only quantization config.
 
@@ -105,7 +106,7 @@ def __init__(
             double_quant_sym (bool): Indicates whether double_quant scale are symmetric, default is True.
             double_quant_group_size (int): Size of double_quant groups, default is 32.
         """
-        super().__init__()
+        super().__init__(white_list=white_list)
         self.weight_bits = weight_bits
         self.weight_dtype = weight_dtype
         self.weight_group_size = weight_group_size
@@ -119,6 +120,7 @@ def __init__(
         self.double_quant_dtype = double_quant_dtype
         self.double_quant_sym = double_quant_sym
         self.double_quant_group_size = double_quant_group_size
+        self._post_init()
 
     def to_dict(self):
         return super().to_dict(params_list=self.params_list, operator2str=operator2str)
@@ -220,12 +222,13 @@ def __init__(
         double_quant_bits: int = 8,
         double_quant_sym: bool = True,
         double_quant_group_size: int = 256,
+        white_list: Optional[List[OP_NAME_OR_MODULE_TYPE]] = DEFAULT_WHITE_LIST,
     ):
         """Init GPTQ config.
 
         Args:
         """
-        super().__init__()
+        super().__init__(white_list=white_list)
         self.weight_dtype = weight_dtype
         self.weight_bits = weight_bits
         self.weight_group_size = weight_group_size
@@ -248,6 +251,7 @@ def __init__(
         self.double_quant_dtype = double_quant_dtype
         self.double_quant_sym = double_quant_sym
         self.double_quant_group_size = double_quant_group_size
+        self._post_init()
 
     def to_dict(self):
         return super().to_dict(params_list=self.params_list, operator2str=operator2str)
@@ -281,65 +285,6 @@ def get_default_gptq_config() -> GPTQConfig:
     return GPTQConfig()
 
 
-######################## Dummy Config ###############################
-# TODO (Yi) remove it after finishing the GPTQ config
-@register_config(framework_name=FRAMEWORK_NAME, algo_name=DUMMY_CONFIG)
-class DummyConfig(BaseConfig):
-    """Config class for round-to-nearest weight-only quantization."""
-
-    supported_configs: List[OperatorConfig] = []
-    params_list = ["act_dtype", "weight_dtype", "dummy_attr"]
-    name = DUMMY_CONFIG
-
-    def __init__(
-        self,
-        weight_dtype: str = "int",
-        act_dtype: str = "fp32",
-        dummy_attr: int = 0,
-    ):
-        """Init RTN weight-only quantization config.
-
-        Args:
-            act_dtype (str): Data type for activations, default is "fp32".
-            weight_dtype (str): Data type for weights, default is "int".
-            dummy_attr (int): Dummy attribute, default is 0.
-        """
-        super().__init__()
-        self.act_dtype = act_dtype
-        self.weight_dtype = weight_dtype
-        self.dummy_attr = dummy_attr
-
-    def to_dict(self):
-        return super().to_dict(params_list=self.params_list, operator2str=operator2str)
-
-    @classmethod
-    def from_dict(cls, config_dict):
-        return super(DummyConfig, cls).from_dict(config_dict=config_dict, str2operator=str2operator)
-
-    @classmethod
-    def register_supported_configs(cls) -> List[OperatorConfig]:
-        supported_configs = []
-        linear_dummy_config = DummyConfig(
-            act_dtype=["fp32"],
-            weight_dtype=["int4", "int8"],
-            dummy_attr=[1, 2, 3],
-        )
-        operators = [torch.nn.Linear, torch.nn.functional.linear]
-        supported_configs.append(
-            OperatorConfig(config=linear_dummy_config, operators=operators, backend=Backend.DEFAULT)
-        )
-        cls.supported_configs = supported_configs
-
-
-def get_default_dummy_config() -> DummyConfig:
-    """Generate the default dummy config.
-
-    Returns:
-        the default dummy config.
-    """
-    return DummyConfig()
-
-
 ##################### Algo Configs End ###################################