quantization_config.py Source File

quantization_config.py Source File#

Mobilint SDK qb Compiler: quantization_config.py Source File
Mobilint SDK qb Compiler v1.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 = ["max", "maxPercentile", "histogram"]
    search_type_list = ["percentile", "mse", "kl"]
 
    DEFAULTS = {
        "quantization_method": 1,
        "quantization_output": 0,
        "quantization_mode": 1,  # maxPercentile is now index 1
        # maxPercentile config
        "percentile": 0.9999,
        "topk_ratio": 0.01,
        "max_each": 128,
        "max_total": 65536,
        # histogram config
        "hist_percentile": 0.9999,
        "hist_use_gpu": True,
        "hist_num_bins": 2048,
        "hist_num_samples": 1024,
        "hist_buffer_size": 1024,
        "hist_min_bin_width": 0.0001,
        "hist_search_percentile_min": 0.9999,
        "hist_search_percentile_max": 0.999999,
        "hist_num_search": 1024,
        "hist_search_type": 0,  # 0=percentile, 1=mse, 2=kl
        # common
        "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",
        # maxPercentile
        "percentile": "percentile",
        "topk_ratio": "topKRatio",
        "max_each": "maxEach",
        "max_total": "maxTotal",
        # histogram
        "hist_percentile": "percentile",
        "hist_use_gpu": "useGPU",
        "hist_num_bins": "numBins",
        "hist_num_samples": "numSamples",
        "hist_buffer_size": "bufferSize",
        "hist_min_bin_width": "minBinWidth",
        "hist_search_percentile_min": "searchPercentileMin",
        "hist_search_percentile_max": "searchPercentileMax",
        "hist_num_search": "numSearch",
        "hist_search_type": "searchType",
        # common
        "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, 2],  # 0=max, 1=maxPercentile, 2=histogram
        # maxPercentile
        "percentile": [0, 1],
        "topk_ratio": [0, 1],
        "max_each": [0, None],
        "max_total": [0, None],
        # histogram
        "hist_percentile": [0, 1],
        "hist_num_bins": [1, None],
        "hist_num_samples": [1, None],
        "hist_buffer_size": [1, None],
        "hist_min_bin_width": [0, None],
        "hist_search_percentile_min": [0, 1],
        "hist_search_percentile_max": [0, 1],
        "hist_num_search": [1, None],
        "hist_search_type": [0, 2],  # 0=percentile, 1=mse, 2=kl
        # common
        "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,
        # maxPercentile config
        percentile: numbers.Number,
        topk_ratio: numbers.Number,
        max_each: int,
        max_total: int,
        # histogram config
        hist_percentile: numbers.Number,
        hist_use_gpu: bool,
        hist_num_bins: int,
        hist_num_samples: int,
        hist_buffer_size: int,
        hist_min_bin_width: numbers.Number,
        hist_search_percentile_min: numbers.Number,
        hist_search_percentile_max: numbers.Number,
        hist_num_search: int,
        hist_search_type: int,
        # common
        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: max.<br>
            1: maxPercentile.<br>
            2: histogram.<br>
        @param percentile float. Percentile value for maxPercentile mode.
        @param topk_ratio float. Top-k ratio used in maxPercentile mode.
        @param max_each int. Maximum number of samples processed per iteration (maxPercentile).
        @param max_total int. Total maximum number of samples (maxPercentile).
        @param hist_percentile float. Percentile for histogram mode (when searchType is percentile).
        @param hist_use_gpu bool. Whether to use GPU for histogram computation.
        @param hist_num_bins int. Number of histogram bins.
        @param hist_num_samples int. Number of representative samples for metric computation.
        @param hist_buffer_size int. Buffer size for batch histogram computation.
        @param hist_min_bin_width float. Minimum bin width for histogram.
        @param hist_search_percentile_min float. Minimum search percentile for MSE/KL search.
        @param hist_search_percentile_max float. Maximum search percentile for MSE/KL search.
        @param hist_num_search int. Number of search iterations for MSE/KL search.
        @param hist_search_type int. Search type: 0=percentile, 1=mse, 2=kl.
        @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.
        @param optional bool. Indicates whether this config is optional.
        """
        super().__init__(optional)
        self.quantization_method = quantization_method
        self.quantization_output = quantization_output
        self.quantization_mode = quantization_mode
        # maxPercentile config
        self.percentile = percentile
        self.topk_ratio = topk_ratio
        self.max_each = max_each
        self.max_total = max_total
        # histogram config
        self.hist_percentile = hist_percentile
        self.hist_use_gpu = hist_use_gpu
        self.hist_num_bins = hist_num_bins
        self.hist_num_samples = hist_num_samples
        self.hist_buffer_size = hist_buffer_size
        self.hist_min_bin_width = hist_min_bin_width
        self.hist_search_percentile_min = hist_search_percentile_min
        self.hist_search_percentile_max = hist_search_percentile_max
        self.hist_num_search = hist_num_search
        self.hist_search_type = hist_search_type
        # common
        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"]
 
        # maxPercentile config (NEW structure)
        if "maxPercentile" in config_dict:
            mp = config_dict["maxPercentile"]
            if "percentile" in mp:
                params["percentile"] = mp["percentile"]
            if "topKRatio" in mp:
                params["topk_ratio"] = mp["topKRatio"]
            if "maxEach" in mp:
                params["max_each"] = mp["maxEach"]
            if "maxTotal" in mp:
                params["max_total"] = mp["maxTotal"]
 
        # histogram config (NEW structure)
        if "histogram" in config_dict:
            h = config_dict["histogram"]
            if "percentile" in h:
                params["hist_percentile"] = h["percentile"]
            if "useGPU" in h:
                params["hist_use_gpu"] = h["useGPU"]
            if "numBins" in h:
                params["hist_num_bins"] = h["numBins"]
            if "numSamples" in h:
                params["hist_num_samples"] = h["numSamples"]
            if "bufferSize" in h:
                params["hist_buffer_size"] = h["bufferSize"]
            if "minBinWidth" in h:
                params["hist_min_bin_width"] = h["minBinWidth"]
            if "searchPercentileMin" in h:
                params["hist_search_percentile_min"] = h["searchPercentileMin"]
            if "searchPercentileMax" in h:
                params["hist_search_percentile_max"] = h["searchPercentileMax"]
            if "numSearch" in h:
                params["hist_num_search"] = h["numSearch"]
            if "searchType" in h:
                # Handle both string and int
                st = h["searchType"]
                if isinstance(st, str):
                    params["hist_search_type"] = cls.search_type_list.index(st.lower())
                else:
                    params["hist_search_type"] = st
 
        # 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,
            "maxPercentile": {
                "percentile": self.percentile,
                "topKRatio": self.topk_ratio,
                "maxEach": self.max_each,
                "maxTotal": self.max_total,
            },
            "histogram": {
                "percentile": self.hist_percentile,
                "useGPU": self.hist_use_gpu,
                "numBins": self.hist_num_bins,
                "numSamples": self.hist_num_samples,
                "bufferSize": self.hist_buffer_size,
                "minBinWidth": self.hist_min_bin_width,
                "searchPercentileMin": self.hist_search_percentile_min,
                "searchPercentileMax": self.hist_search_percentile_max,
                "numSearch": self.hist_num_search,
                "searchType": self.search_type_list[self.hist_search_type],
            },
            "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),
            optional=cls.optional,
        )
 
    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
    )
 
 
# @}