quantization_config.py Source File

quantization_config.py Source File#

Mobilint SDK qb Compiler: quantization_config.py Source File
Mobilint SDK qb Compiler v0.12.0.0
MCS002-KR
import numbers
from typing import Dict, List, Optional
from typeguard import typechecked
from . import ConfigABC
 
 
 
class CalibrationConfig(ConfigABC):
    """
    @brief Configuration for calibration during quantization.
 
    @details Defines calibration and quantization parameterization used to derive activation/weight scales and related
    statistics during quantized compilation.
    """
 
    optional = False
    method_list = ["WChALayer", "WChAMulti", "WChALayerZeropoint", "WChAMultiZeropoint"]
    output_list = ["Layer", "Ch", "Sigmoid"]
    mode_list = [
        "percentile",
        "max",
        "maxPercentile",
        "fastPercentile",
        "histogramKL",
        "histogramMSE",
    ]
 
    DEFAULTS = {
        "quantization_method": 1,
        "quantization_output": 0,
        "quantization_mode": 2,
        "percentile": 0.9999,
        "topk_ratio": 0.01,
        "max_each": 128,
        "max_total": 65536,
        "kernel_size": 9,
        "stack_size": 32768,
        "hist_num_cluster": 4,
        "hist_search_percentile_min": 0.9999,
        "hist_epsilon": 0.01,
        "hist_adjust_coeff": 0.5,
        "act_scale_min": 0.0005,
        "weight_scale_min": 1e-6,
        "min_clip_ratio": -1,
        "layer_overrides": {"actScaleMin": {"0.0005": []}},
    }
 
    py_to_cpp_name = {
        "quantization_method": "method",
        "quantization_output": "output",
        "quantization_mode": "mode",
        "percentile": "percentile",
        "topk_ratio": "topKRatio",
        "max_each": "maxEach",
        "max_total": "maxTotal",
        "kernel_size": "kernelSize",
        "stack_size": "stackSize",
        "hist_num_cluster": "numCluster",
        "hist_search_percentile_min": "searchPercentileMin",
        "hist_epsilon": "epsilon",
        "hist_adjust_coeff": "adjustCoeff",
        "act_scale_min": "actScaleMin",
        "weight_scale_min": "weightScaleMin",
        "min_clip_ratio": "minClipRatio",
    }
 
    min_max_check = {
        "quantization_method": [0, 3],
        "quantization_output": [0, 2],
        "quantization_mode": [0, 5],
        "percentile": [0, 1],
        "topk_ratio": [0, 1],
        "max_each": [0, None],
        "max_total": [0, None],
        "kernel_size": [0, None],
        "stack_size": [0, None],
        "hist_num_cluster": [0, None],
        "hist_search_percentile_min": [0, 1],
        "hist_epsilon": [0, 1],
        "hist_adjust_coeff": [0, None],
        "act_scale_min": [0, 1],
        "weight_scale_min": [0, 1],
        "min_clip_ratio": [-1, 1],
    }
 
    def __init__(
        self,
        quantization_method: int,
        quantization_output: int,
        quantization_mode: int,
        percentile: numbers.Number,
        topk_ratio: numbers.Number,
        max_each: int,
        max_total: int,
        kernel_size: int,
        stack_size: int,
        hist_num_cluster: int,
        hist_search_percentile_min: numbers.Number,
        hist_epsilon: numbers.Number,
        hist_adjust_coeff: numbers.Number,
        act_scale_min: numbers.Number,
        weight_scale_min: numbers.Number,
        min_clip_ratio: numbers.Number,
        layer_overrides: Dict,
        optional: bool = optional,
    ):
        """
        @brief Initialize the CalibrationConfig.
 
        @param quantization_method int. Calibration method index:<br>
            0: WChALayer - Weight per-channel, Activation per-layer.<br>
            1: WChAMulti - Weight per-channel, Activation multi-layer.<br>
            2: WChALayerZeropoint - Weight per-channel, Activation per-layer with zeropoint.<br>
            3: WChAMultiZeropoint - Weight per-channel, Activation multi-layer with zeropoint.<br>
        @param quantization_output int. Output quantization type index:<br>
            0: Layer - Per-layer quantization.<br>
            1: Ch - Per-channel quantization.<br>
            2: Sigmoid - Sigmoid-based quantization.<br>
        @param quantization_mode int. Quantization mode index:<br>
            0: percentile.<br>
            1: max.<br>
            2: maxPercentile.<br>
            3: fastPercentile.<br>
            4: histogramKL.<br>
            5: histogramMSE.<br>
        @param percentile float. Percentile value used when @p quantization_mode is 0 (percentile).
        @param topk_ratio float. Top-k ratio used when @p quantization_mode is 2 (maxPercentile).
        @param max_each int. Maximum number of samples processed per iteration.
        @param max_total int. Total maximum number of samples processed across all iterations.
        @param kernel_size int. Kernel size used for fast distribution estimation (typically for @p quantization_mode 3).
        @param stack_size int. Stack size used for fast distribution estimation (typically for @p quantization_mode 3).
        @param hist_num_cluster int. Number of clusters/bins used in histogram-based calibration modes.
        @param hist_search_percentile_min float. Minimum search percentile for histogram-based calibration.
        @param hist_epsilon float. Numerical epsilon used in histogram-based calibration to avoid instability.
        @param hist_adjust_coeff float. Adjustment coefficient applied in histogram-based calibration.
        @param act_scale_min float. Minimum allowed activation scale (lower bound clamp).
        @param weight_scale_min float. Minimum allowed weight scale (lower bound clamp).
        @param min_clip_ratio float. Minimum clip ratio constraint applied during calibration.
        @param layer_overrides Dict. Layer-specific override settings. Keys are layer identifiers; values override one or more
        calibration parameters for the specified layer.
        @param optional bool. Indicates whether this config is optional. Use the default value defined as a class variable.
        """
        super().__init__(optional)
        self.quantization_method = quantization_method
        self.quantization_output = quantization_output
        self.quantization_mode = quantization_mode
        self.percentile = percentile
        self.topk_ratio = topk_ratio
        self.max_each = max_each
        self.max_total = max_total
        self.kernel_size = kernel_size
        self.stack_size = stack_size
        self.hist_num_cluster = hist_num_cluster
        self.hist_search_percentile_min = hist_search_percentile_min
        self.hist_epsilon = hist_epsilon
        self.hist_adjust_coeff = hist_adjust_coeff
        self.act_scale_min = act_scale_min
        self.weight_scale_min = weight_scale_min
        self.min_clip_ratio = min_clip_ratio
        self.layer_overrides = layer_overrides
        self.check_valid(except_list=["optional", "layer_overrides"])
 
    @classmethod
    def default_config(cls):
        return cls(**cls.DEFAULTS, optional=cls.optional)
 
    @classmethod
    def from_kwargs(cls, **kwargs):
        params = cls.DEFAULTS.copy()
        updated_kwargs = dict()
        for key in cls.DEFAULTS.keys():
            if key in kwargs:
                updated_kwargs[key] = kwargs[key]
                params[key] = kwargs[key]
 
        try:
            return cls(**params, optional=cls.optional)
        except:
            raise ValueError(
                f"Please review the configuration kwargs passed to {cls.__name__}. These kwargs were updated and will be used: {updated_kwargs}"
            )
 
    @classmethod
    def from_dict(cls, config_dict: dict):
        """Create CalibrationConfig from dictionary (JSON structure)"""
        params = cls.DEFAULTS.copy()
 
        # Top-level parameters
        if "method" in config_dict:
            params["quantization_method"] = config_dict["method"]
        if "output" in config_dict:
            params["quantization_output"] = config_dict["output"]
        if "mode" in config_dict:
            params["quantization_mode"] = config_dict["mode"]
 
        # Percentile config
        if "percentile" in config_dict:
            percentile_config = config_dict["percentile"]
            if "percentile" in percentile_config:
                params["percentile"] = percentile_config["percentile"]
            if "max" in percentile_config:
                max_config = percentile_config["max"]
                if "topKRatio" in max_config:
                    params["topk_ratio"] = max_config["topKRatio"]
                if "maxEach" in max_config:
                    params["max_each"] = max_config["maxEach"]
                if "maxTotal" in max_config:
                    params["max_total"] = max_config["maxTotal"]
 
        # Fast distribution config
        if "fastDist" in config_dict:
            fast_dist = config_dict["fastDist"]
            if "kernelSize" in fast_dist:
                params["kernel_size"] = fast_dist["kernelSize"]
            if "stackSize" in fast_dist:
                params["stack_size"] = fast_dist["stackSize"]
 
        # Histogram config
        if "histogram" in config_dict:
            histogram = config_dict["histogram"]
            if "numCluster" in histogram:
                params["hist_num_cluster"] = histogram["numCluster"]
            if "searchPercentileMin" in histogram:
                params["hist_search_percentile_min"] = histogram["searchPercentileMin"]
            if "epsilon" in histogram:
                params["hist_epsilon"] = histogram["epsilon"]
            if "adjustCoeff" in histogram:
                params["hist_adjust_coeff"] = histogram["adjustCoeff"]
 
        # Scale and clip parameters
        if "actScaleMin" in config_dict:
            params["act_scale_min"] = config_dict["actScaleMin"]
        if "weightScaleMin" in config_dict:
            params["weight_scale_min"] = config_dict["weightScaleMin"]
        if "minclipRatio" in config_dict:
            params["min_clip_ratio"] = config_dict["minclipRatio"]
        if "layerOverrides" in config_dict:
            params["layer_overrides"] = config_dict["layerOverrides"]
 
        return cls(**params)
 
    def to_dict(self):
        return {
            "minclipRatio": self.min_clip_ratio,
            "methodList": self.method_list,
            "method": self.quantization_method,
            "outputList": self.output_list,
            "output": self.quantization_output,
            "modeList": self.mode_list,
            "mode": self.quantization_mode,
            "percentile": {
                "percentile": self.percentile,
                "max": {
                    "topKRatio": self.topk_ratio,
                    "maxEach": self.max_each,
                    "maxTotal": self.max_total,
                },
            },
            "max": {},
            "fastDist": {
                "sizeCali": 100,
                "kernelSize": self.kernel_size,
                "stackSize": self.stack_size,
            },
            "histogram": {
                "numCluster": self.hist_num_cluster,
                "searchPercentileMin": self.hist_search_percentile_min,
                "epsilon": self.hist_epsilon,
                "adjustCoeff": self.hist_adjust_coeff,
            },
            "kurtosis": {
                "percentileMin": 0.9999,
                "percentileMax": 0.999999,
                "kurtosisMin": 3.0,
                "kurtosisMax": 1000.0,
                "convexity": 1.0,
            },
            "actScaleMin": self.act_scale_min,
            "weightScaleMin": self.weight_scale_min,
            "layerOverrides": self.layer_overrides,
        }
 
 
class BitConfig(ConfigABC):
    """
    @brief Configuration for bit precision.
    @details Defines bit-width parameterization for activations and weights used in mixed-precision quantization (e.g., attention and FFN components).
    """
 
    optional = False
 
    DEFAULTS = {
        "query_act_bits": 8,
        "key_act_bits": 8,
        "value_act_bits": 8,
        "output_act_bits": 16,
        "ffn_act_bits": 16,
        "head_act_bits": 8,
        "query_weight_bits": 8,
        "key_weight_bits": 8,
        "value_weight_bits": 8,
        "output_weight_bits": 8,
        "ffn_weight_bits": 8,
        "head_weight_bits": 8,
        "mixed_precision_apply": False,
        "save_path": "",
        "load_path": "",
        "activation_16bits": [],
        "weight_16bits": [],
    }
 
    py_to_cpp_name = {
        "query_act_bits": "queryActBits",
        "key_act_bits": "keyActBits",
        "value_act_bits": "valueActBits",
        "output_act_bits": "outputActBits",
        "ffn_act_bits": "ffnActBits",
        "head_act_bits": "headActBits",
        "query_weight_bits": "queryWeightBits",
        "key_weight_bits": "keyWeightBits",
        "value_weight_bits": "valueWeightBits",
        "output_weight_bits": "outputWeightBits",
        "ffn_weight_bits": "ffnWeightBits",
        "head_weight_bits": "headWeightBits",
        "mixed_precision_apply": "mixedPrecisionApply",
        "save_path": "savePath",
        "load_path": "loadPath",
    }
 
    # Check that all bits values are positive
    min_max_check = {
        "query_act_bits": [1, None],
        "key_act_bits": [1, None],
        "value_act_bits": [1, None],
        "output_act_bits": [1, None],
        "ffn_act_bits": [1, None],
        "head_act_bits": [1, None],
        "query_weight_bits": [1, None],
        "key_weight_bits": [1, None],
        "value_weight_bits": [1, None],
        "output_weight_bits": [1, None],
        "ffn_weight_bits": [1, None],
        "head_weight_bits": [1, None],
    }
 
    def __init__(
        self,
        query_act_bits: int,
        key_act_bits: int,
        value_act_bits: int,
        output_act_bits: int,
        ffn_act_bits: int,
        head_act_bits: int,
        query_weight_bits: int,
        key_weight_bits: int,
        value_weight_bits: int,
        output_weight_bits: int,
        ffn_weight_bits: int,
        head_weight_bits: int,
        mixed_precision_apply: bool,
        save_path: str,
        load_path: str,
        activation_16bits: List[str],
        weight_16bits: List[str],
        optional: bool = optional,
    ):
        """
        @brief Initialize the BitConfig.
 
        @param query_act_bits int. Query activation bit-width.
        @param key_act_bits int. Key activation bit-width.
        @param value_act_bits int. Value activation bit-width.
        @param output_act_bits int. Output activation bit-width.
        @param ffn_act_bits int. FFN activation bit-width.
        @param head_act_bits int. Head activation bit-width.
        @param query_weight_bits int. Query weight bit-width.
        @param key_weight_bits int. Key weight bit-width.
        @param value_weight_bits int. Value weight bit-width.
        @param output_weight_bits int. Output weight bit-width.
        @param ffn_weight_bits int. FFN weight bit-width.
        @param head_weight_bits int. Head weight bit-width.
        @param mixed_precision_apply bool. If true, apply mixed-precision according to the specified bit-widths.
        @param save_path str. (optional) Path to save the bit allocation. If empty, the allocation is not saved.
        @param load_path str. (optional) Path to load the bit allocation. If empty, allocation is not loaded.
        @param activation_16bits List[str]. (optional) Layer names to force 16-bit activations.
        @param weight_16bits List[str]. (optional) Layer names to force 16-bit weights.
        """
        super().__init__(optional)
        self.query_act_bits = query_act_bits
        self.key_act_bits = key_act_bits
        self.value_act_bits = value_act_bits
        self.output_act_bits = output_act_bits
        self.ffn_act_bits = ffn_act_bits
        self.head_act_bits = head_act_bits
        self.query_weight_bits = query_weight_bits
        self.key_weight_bits = key_weight_bits
        self.value_weight_bits = value_weight_bits
        self.output_weight_bits = output_weight_bits
        self.ffn_weight_bits = ffn_weight_bits
        self.head_weight_bits = head_weight_bits
        self.mixed_precision_apply = mixed_precision_apply
        self.save_path = save_path
        self.load_path = load_path
        self.activation_16bits = activation_16bits
        self.weight_16bits = weight_16bits
        self.check_valid(except_list=["optional", "activation_16bits", "weight_16bits"])
 
    @classmethod
    def default_config(cls):
        return cls(**cls.DEFAULTS, optional=cls.optional)
 
    @classmethod
    def from_kwargs(cls, **kwargs):
        params = cls.DEFAULTS.copy()
        updated_kwargs = dict()
        for key in cls.DEFAULTS.keys():
            if key in kwargs:
                updated_kwargs[key] = kwargs[key]
                params[key] = kwargs[key]
 
        try:
            return cls(**params, optional=cls.optional)
        except:
            raise ValueError(
                f"Please review the configuration kwargs passed to {cls.__name__}. These kwargs were updated and will be used: {updated_kwargs}"
            )
 
    @classmethod
    def from_dict(cls, config_dict: dict):
        """Create BitConfig from dictionary (JSON structure)"""
        params = cls.DEFAULTS.copy()
 
        # Transformer config
        if "transformer" in config_dict:
            transformer = config_dict["transformer"]
 
            # Activation bits
            if "activation" in transformer:
                activation = transformer["activation"]
                if "query" in activation:
                    params["query_act_bits"] = activation["query"]
                if "key" in activation:
                    params["key_act_bits"] = activation["key"]
                if "value" in activation:
                    params["value_act_bits"] = activation["value"]
                if "output" in activation:
                    params["output_act_bits"] = activation["output"]
                if "ffn" in activation:
                    params["ffn_act_bits"] = activation["ffn"]
                if "head" in activation:
                    params["head_act_bits"] = activation["head"]
 
            # Weight bits
            if "weight" in transformer:
                weight = transformer["weight"]
                if "query" in weight:
                    params["query_weight_bits"] = weight["query"]
                if "key" in weight:
                    params["key_weight_bits"] = weight["key"]
                if "value" in weight:
                    params["value_weight_bits"] = weight["value"]
                if "output" in weight:
                    params["output_weight_bits"] = weight["output"]
                if "ffn" in weight:
                    params["ffn_weight_bits"] = weight["ffn"]
                if "head" in weight:
                    params["head_weight_bits"] = weight["head"]
 
            # Mixed precision
            if "mixedPrecision" in transformer:
                mixed_precision = transformer["mixedPrecision"]
                if "apply" in mixed_precision:
                    params["mixed_precision_apply"] = mixed_precision["apply"]
 
        # Save info
        if "saveInfo" in config_dict:
            save_info = config_dict["saveInfo"]
            if "savePath" in save_info:
                params["save_path"] = save_info["savePath"]
            if "loadPath" in save_info:
                params["load_path"] = save_info["loadPath"]
 
        # Layer overrides
        if "layerOverrides" in config_dict:
            layer_overrides = config_dict["layerOverrides"]
            if "activation16Bits" in layer_overrides:
                params["activation_16bits"] = layer_overrides["activation16Bits"]
            if "weight16Bits" in layer_overrides:
                params["weight_16bits"] = layer_overrides["weight16Bits"]
 
        return cls(**params)
 
    def to_dict(self):
        return {
            "transformer": {
                "activation": {
                    "query": self.query_act_bits,
                    "key": self.key_act_bits,
                    "value": self.value_act_bits,
                    "output": self.output_act_bits,
                    "ffn": self.ffn_act_bits,
                    "head": self.head_act_bits,
                },
                "weight": {
                    "query": self.query_weight_bits,
                    "key": self.key_weight_bits,
                    "value": self.value_weight_bits,
                    "output": self.output_weight_bits,
                    "ffn": self.ffn_weight_bits,
                    "head": self.head_weight_bits,
                },
                "mixedPrecision": {
                    "apply": self.mixed_precision_apply,
                    "typeWise": True,
                    "prune": 0,
                    "bit_2": 0,
                    "bit_4": 0,
                    "bit_8": 1,
                    "importanceThreshold_low": -1,
                    "importanceThreshold_high": -1,
                },
            },
            "saveInfo": {
                "savePath": self.save_path,
                "loadPath": self.load_path,
            },
            "layerOverrides": {
                "activation16Bits": self.activation_16bits,
                "weight16Bits": self.weight_16bits,
            },
        }
 
 
class QuantizationConfig(ConfigABC):
    """
    @brief Unified quantization configuration.
 
    @details Groups calibration and bit-precision configurations into a single configuration.
    """
 
    optional = False
    py_to_cpp_name = {
        "calibration": "calibration",
        "bit": "bit",
    }
 
    min_max_check = {}
 
    @typechecked
    def __init__(
        self,
        optional: bool,
        calibration: Optional[CalibrationConfig] = None,
        bit: Optional[BitConfig] = None,
    ):
        """
        @brief Initialize the QuantizationConfig.
        @param calibration CalibrationConfig. Calibration configuration.
        @param bit BitConfig. Bit precision configuration.
        """
        super().__init__(optional)
        self.calibration = calibration or CalibrationConfig.default_config()
        self.bit = bit or BitConfig.default_config()
 
    @classmethod
    def default_config(cls):
        return QuantizationConfig(
            calibration=CalibrationConfig.default_config(),
            bit=BitConfig.default_config(),
            optional=cls.optional,
        )
 
    @classmethod
    def from_kwargs(cls, **kwargs):
        return QuantizationConfig(
            calibration=CalibrationConfig.from_kwargs(**kwargs),
            bit=BitConfig.from_kwargs(**kwargs),
            optional=cls.optional,
        )
 
    @classmethod
    def from_dict(cls, config_dict: dict):
        calibration_dict = config_dict.get("calibration", {})
        bit_dict = config_dict.get("bit", {})
 
        return cls(
            calibration=CalibrationConfig.from_dict(calibration_dict),
            bit=BitConfig.from_dict(bit_dict),
        )
 
    def to_dict(self):
        return {
            "quantization": {
                "calibration": self.calibration.to_dict(),
                "bit": self.bit.to_dict(),
            }
        }
 
 
def get_calibration_config(**kwargs) -> CalibrationConfig:
    """
    @brief Create CalibrationConfig with partial parameters merged with defaults.
    """
    params = CalibrationConfig.DEFAULTS.copy()
    params.update(kwargs)
    return CalibrationConfig(**params, optional=CalibrationConfig.optional)
 
 
def get_bit_config(**kwargs) -> BitConfig:
    """
    @brief Create BitConfig with partial parameters merged with defaults.
    """
    params = BitConfig.DEFAULTS.copy()
    params.update(kwargs)
    return BitConfig(**params, optional=BitConfig.optional)
 
 
def get_quantization_config(
    calibration: Optional[CalibrationConfig] = None,
    bit: Optional[BitConfig] = None,
) -> QuantizationConfig:
    """
    @brief Create QuantizationConfig with partial parameters merged with defaults.
    """
    return QuantizationConfig(
        calibration=calibration, bit=bit, optional=QuantizationConfig.optional
    )
 
 
# @}