equivalent_transformation_config.py Source File

equivalent_transformation_config.py Source File#

Mobilint SDK qb Compiler: equivalent_transformation_config.py Source File
Mobilint SDK qb Compiler v1.0
MCS002-EN
from typing import Dict, List, Optional
from typeguard import typechecked
from . import ConfigABC
from .config_utils import value_range_check
 
 
 
class EquivalentTransformationConfig(ConfigABC):
    """
    @brief Configuration for equivalent transformation techniques.
 
    @details This includes various transformation methods like NormConv, QK, UD, VO,
    FeedForwardMultiLUT, SpinR1, HeadOutChRotation, SpinR2, QKRotation, and FlattenQuant.
    """
 
    optional = False
    py_to_cpp_name = {}
    # Empty py_to_cpp_name since this config uses nested structure in to_dict
 
    # Default values shared between __init__ and from_dict
    DEFAULTS = {
        "seed": 0,
        "apply_hadamard_rotation_matrix": True,
        # NormConv
        "norm_conv_apply": False,
        "norm_conv_learn": False,
        "norm_conv_smoothing_factor": 0.5,
        "norm_conv_min_gamma": 0.0001,
        "norm_conv_max_gamma": 10000.0,
        # QK
        "qk_apply": False,
        "qk_smoothing_factor": 0.5,
        "qk_min_gamma": 0.0001,
        "qk_max_gamma": 10000.0,
        # UD
        "ud_apply": False,
        "ud_learn": False,
        "ud_smoothing_factor": 0.5,
        "ud_min_gamma": 0.0001,
        "ud_max_gamma": 10000.0,
        # VO
        "vo_apply": False,
        "vo_smoothing_factor": 0.5,
        "vo_min_gamma": 0.0001,
        "vo_max_gamma": 10000.0,
        # FeedForwardMultiLUT
        "ff_multi_lut_apply": False,
        "ff_multi_lut_breakpoints": [-8.0, -4.0, 0],
        # SpinR1
        "spin_r1_apply": False,
        "spin_r1_matrix_path": "",
        # HeadOutChRotation
        "head_out_ch_rotation_apply": False,
        "head_out_ch_rotation_matrix_path": "",
        # SpinR2
        "spin_r2_apply": False,
        "spin_r2_learn": False,
        "spin_r2_matrix_path": "",
        # QKRotation
        "qk_rotation_apply": False,
        "qk_rotation_matrix_path": "",
        # FlattenQuant
        "flatten_quant_apply": False,
        "flatten_quant_learn": False,
        "flatten_quant_apply_threshold": 0.33,
        "flatten_quant_max_overhead": 0.05,
        # OptimizeFFN
        "optimize_ffn_apply": False,
        "optimize_ffn_ch_per_ffn": -1,
    }
 
    # Min/max range checks
    min_max_check = {
        "norm_conv_smoothing_factor": [0.0, 1.0],
        "norm_conv_min_gamma": [0.0, None],
        "norm_conv_max_gamma": [0.0, None],
        "qk_smoothing_factor": [0.0, 1.0],
        "qk_min_gamma": [0.0, None],
        "qk_max_gamma": [0.0, None],
        "ud_smoothing_factor": [0.0, 1.0],
        "ud_min_gamma": [0.0, None],
        "ud_max_gamma": [0.0, None],
        "vo_smoothing_factor": [0.0, 1.0],
        "vo_min_gamma": [0.0, None],
        "vo_max_gamma": [0.0, None],
        "flatten_quant_apply_threshold": [0.0, 1.0],
        "flatten_quant_max_overhead": [0.0, None],
    }
 
    def __init__(
        self,
        seed: int,
        apply_hadamard_rotation_matrix: bool,
        # NormConv
        norm_conv_apply: bool,
        norm_conv_learn: bool,
        norm_conv_smoothing_factor: float,
        norm_conv_min_gamma: float,
        norm_conv_max_gamma: float,
        # QK
        qk_apply: bool,
        qk_smoothing_factor: float,
        qk_min_gamma: float,
        qk_max_gamma: float,
        # UD
        ud_apply: bool,
        ud_learn: bool,
        ud_smoothing_factor: float,
        ud_min_gamma: float,
        ud_max_gamma: float,
        # VO
        vo_apply: bool,
        vo_smoothing_factor: float,
        vo_min_gamma: float,
        vo_max_gamma: float,
        # FeedForwardMultiLUT
        ff_multi_lut_apply: bool,
        ff_multi_lut_breakpoints: List[float],
        # SpinR1
        spin_r1_apply: bool,
        spin_r1_matrix_path: str,
        # HeadOutChRotation
        head_out_ch_rotation_apply: bool,
        head_out_ch_rotation_matrix_path: str,
        # SpinR2
        spin_r2_apply: bool,
        spin_r2_learn: bool,
        spin_r2_matrix_path: str,
        # QKRotation
        qk_rotation_apply: bool,
        qk_rotation_matrix_path: str,
        # FlattenQuant
        flatten_quant_apply: bool,
        flatten_quant_learn: bool,
        flatten_quant_apply_threshold: float,
        flatten_quant_max_overhead: float,
        # OptimizeFFN
        optimize_ffn_apply: bool,
        optimize_ffn_ch_per_ffn: int,
        optional: bool = optional,
    ):
        """
        @brief Initialize the EquivalentTransformationConfig.
        @param seed int. Random seed used for equivalent transformation operations. Default is 0.
        @param apply_hadamard_rotation_matrix bool. If True, apply the Hadamard rotation matrix where applicable. Default is True.
 
        # NormConv
        @param norm_conv_apply bool. If True, enable NormConv transformation. Default is False.
        @param norm_conv_learn bool. If True, learn NormConv parameters during calibration/optimization. Default is False.
        @param norm_conv_smoothing_factor float. Smoothing factor used in NormConv. Default is 0.5.
        @param norm_conv_min_gamma float. Minimum gamma value (lower bound clamp) for NormConv. Default is 0.0001.
        @param norm_conv_max_gamma float. Maximum gamma value (upper bound clamp) for NormConv. Default is 10000.0.
 
        # QK
        @param qk_apply bool. If True, enable QK transformation. Default is False.
        @param qk_smoothing_factor float. Smoothing factor used in QK. Default is 0.5.
        @param qk_min_gamma float. Minimum gamma value (lower bound clamp) for QK. Default is 0.0001.
        @param qk_max_gamma float. Maximum gamma value (upper bound clamp) for QK. Default is 10000.0.
 
        # UD
        @param ud_apply bool. If True, enable UD transformation. Default is False.
        @param ud_learn bool. If True, learn UD parameters during calibration/optimization. Default is False.
        @param ud_smoothing_factor float. Smoothing factor used in UD. Default is 0.5.
        @param ud_min_gamma float. Minimum gamma value (lower bound clamp) for UD. Default is 0.0001.
        @param ud_max_gamma float. Maximum gamma value (upper bound clamp) for UD. Default is 10000.0.
 
        # VO
        @param vo_apply bool. If True, enable VO transformation. Default is False.
        @param vo_smoothing_factor float. Smoothing factor used in VO. Default is 0.5.
        @param vo_min_gamma float. Minimum gamma value (lower bound clamp) for VO. Default is 0.0001.
        @param vo_max_gamma float. Maximum gamma value (upper bound clamp) for VO. Default is 10000.0.
 
 
        # FeedForwardMultiLUT
        @param ff_multi_lut_apply bool. If True, enable FeedForward Multi-LUT. Default is False.
        @param ff_multi_lut_breakpoints List[float]. Breakpoints used for FeedForward Multi-LUT piecewise regions. Default is [-8.0, -4.0, 0].
 
        # SpinR1
        @param spin_r1_apply bool. If True, enable SpinR1 rotation. Default is False.
        @param spin_r1_matrix_path str. Path to the SpinR1 rotation matrix file. Default is "".
 
        # HeadOutChRotation
        @param head_out_ch_rotation_apply bool. If True, enable HeadOutChRotation. Default is False.
        @param head_out_ch_rotation_matrix_path str. Path to the HeadOutChRotation matrix file. Default is "".
 
        # SpinR2
        @param spin_r2_apply bool. If True, enable SpinR2 rotation. Default is False.
        @param spin_r2_learn bool. If True, learn the SpinR2 rotation matrix/parameters during calibration/optimization. Default is False.
        @param spin_r2_matrix_path str. Path to the SpinR2 rotation matrix file. Default is "".
 
        # QKRotation
        @param qk_rotation_apply bool. If True, enable QKRotation. Default is False.
        @param qk_rotation_matrix_path str. Path to the QKRotation matrix file. Default is "".
 
        # FlattenQuant
        @param flatten_quant_apply bool. If True, enable FlattenQuant. Default is False.
        @param flatten_quant_learn bool. If True, learn FlattenQuant parameters during calibration/optimization. Default is False.
        @param flatten_quant_apply_threshold float. Threshold for applying FlattenQuant. Default is 0.33.
        @param flatten_quant_max_overhead float. Maximum allowed overhead when applying FlattenQuant. Default is 0.05.
 
        # OptimizeFFN
        @param optimize_ffn_apply bool. If True, enable OptimizeFFN. Default is False.
        @param optimize_ffn_ch_per_ffn int. Number of channels per FFN used when OptimizeFFN is enabled. Default is -1.
 
        @param optional bool. Indicates whether this config is optional. Use the default value defined as a class variable.
        """
        super().__init__(optional)
 
        self.seed = seed
        self.apply_hadamard_rotation_matrix = apply_hadamard_rotation_matrix
 
        # NormConv
        self.norm_conv_apply = norm_conv_apply
        self.norm_conv_learn = norm_conv_learn
        self.norm_conv_smoothing_factor = norm_conv_smoothing_factor
        self.norm_conv_min_gamma = norm_conv_min_gamma
        self.norm_conv_max_gamma = norm_conv_max_gamma
 
        # QK
        self.qk_apply = qk_apply
        self.qk_smoothing_factor = qk_smoothing_factor
        self.qk_min_gamma = qk_min_gamma
        self.qk_max_gamma = qk_max_gamma
 
        # UD
        self.ud_apply = ud_apply
        self.ud_learn = ud_learn
        self.ud_smoothing_factor = ud_smoothing_factor
        self.ud_min_gamma = ud_min_gamma
        self.ud_max_gamma = ud_max_gamma
 
        # VO
        self.vo_apply = vo_apply
        self.vo_smoothing_factor = vo_smoothing_factor
        self.vo_min_gamma = vo_min_gamma
        self.vo_max_gamma = vo_max_gamma
 
        # FeedForwardMultiLUT
        self.ff_multi_lut_apply = ff_multi_lut_apply
        self.ff_multi_lut_breakpoints = ff_multi_lut_breakpoints
 
        # SpinR1
        self.spin_r1_apply = spin_r1_apply
        self.spin_r1_matrix_path = spin_r1_matrix_path
 
        # HeadOutChRotation
        self.head_out_ch_rotation_apply = head_out_ch_rotation_apply
        self.head_out_ch_rotation_matrix_path = head_out_ch_rotation_matrix_path
 
        # SpinR2
        self.spin_r2_apply = spin_r2_apply
        self.spin_r2_learn = spin_r2_learn
        self.spin_r2_matrix_path = spin_r2_matrix_path
 
        # QKRotation
        self.qk_rotation_apply = qk_rotation_apply
        self.qk_rotation_matrix_path = qk_rotation_matrix_path
 
        # FlattenQuant
        self.flatten_quant_apply = flatten_quant_apply
        self.flatten_quant_learn = flatten_quant_learn
        self.flatten_quant_apply_threshold = flatten_quant_apply_threshold
        self.flatten_quant_max_overhead = flatten_quant_max_overhead
 
        # OptimizeFFN
        self.optimize_ffn_apply = optimize_ffn_apply
        self.optimize_ffn_ch_per_ffn = optimize_ffn_ch_per_ffn
 
        self.check_valid()
 
    def check_valid(self, except_list=["optional"]):
        """Override to skip py_to_cpp_name check (uses nested structure in to_dict)"""
        self.check_min_max_range()
        self.check_valid_str_configs()
 
    @classmethod
    def default_config(cls):
        return cls(**cls.DEFAULTS, optional=cls.optional)
 
    @classmethod
    def from_dict(cls, config_dict: dict):
        """Create EquivalentTransformationConfig from dictionary (JSON structure)"""
        # Start with defaults and update with provided values
        params = cls.DEFAULTS.copy()
 
        # Top-level parameters
        if "seed" in config_dict:
            params["seed"] = config_dict["seed"]
        if "applyHadamardRotationMatrix" in config_dict:
            params["apply_hadamard_rotation_matrix"] = config_dict[
                "applyHadamardRotationMatrix"
            ]
 
        # NormConv
        if "NormConv" in config_dict:
            norm_conv = config_dict["NormConv"]
            if "apply" in norm_conv:
                params["norm_conv_apply"] = norm_conv["apply"]
            if "learn" in norm_conv:
                params["norm_conv_learn"] = norm_conv["learn"]
            if "smoothingFactor" in norm_conv:
                params["norm_conv_smoothing_factor"] = norm_conv["smoothingFactor"]
            if "minGamma" in norm_conv:
                params["norm_conv_min_gamma"] = norm_conv["minGamma"]
            if "maxGamma" in norm_conv:
                params["norm_conv_max_gamma"] = norm_conv["maxGamma"]
 
        # QK
        if "QK" in config_dict:
            qk = config_dict["QK"]
            if "apply" in qk:
                params["qk_apply"] = qk["apply"]
            if "smoothingFactor" in qk:
                params["qk_smoothing_factor"] = qk["smoothingFactor"]
            if "minGamma" in qk:
                params["qk_min_gamma"] = qk["minGamma"]
            if "maxGamma" in qk:
                params["qk_max_gamma"] = qk["maxGamma"]
 
        # UD
        if "UD" in config_dict:
            ud = config_dict["UD"]
            if "apply" in ud:
                params["ud_apply"] = ud["apply"]
            if "learn" in ud:
                params["ud_learn"] = ud["learn"]
            if "smoothingFactor" in ud:
                params["ud_smoothing_factor"] = ud["smoothingFactor"]
            if "minGamma" in ud:
                params["ud_min_gamma"] = ud["minGamma"]
            if "maxGamma" in ud:
                params["ud_max_gamma"] = ud["maxGamma"]
 
        # VO
        if "VO" in config_dict:
            vo = config_dict["VO"]
            if "apply" in vo:
                params["vo_apply"] = vo["apply"]
            if "smoothingFactor" in vo:
                params["vo_smoothing_factor"] = vo["smoothingFactor"]
            if "minGamma" in vo:
                params["vo_min_gamma"] = vo["minGamma"]
            if "maxGamma" in vo:
                params["vo_max_gamma"] = vo["maxGamma"]
 
        # FeedForwardMultiLUT
        if "FeedForwardMultiLUT" in config_dict:
            ff_multi_lut = config_dict["FeedForwardMultiLUT"]
            if "apply" in ff_multi_lut:
                params["ff_multi_lut_apply"] = ff_multi_lut["apply"]
            if "breakpoints" in ff_multi_lut:
                params["ff_multi_lut_breakpoints"] = ff_multi_lut["breakpoints"]
 
        # SpinR1
        if "SpinR1" in config_dict:
            spin_r1 = config_dict["SpinR1"]
            if "apply" in spin_r1:
                params["spin_r1_apply"] = spin_r1["apply"]
            if "matrixPath" in spin_r1:
                params["spin_r1_matrix_path"] = spin_r1["matrixPath"]
 
        # HeadOutChRotation
        if "HeadOutChRotation" in config_dict:
            head_out_ch = config_dict["HeadOutChRotation"]
            if "apply" in head_out_ch:
                params["head_out_ch_rotation_apply"] = head_out_ch["apply"]
            if "matrixPath" in head_out_ch:
                params["head_out_ch_rotation_matrix_path"] = head_out_ch["matrixPath"]
 
        # SpinR2
        if "SpinR2" in config_dict:
            spin_r2 = config_dict["SpinR2"]
            if "apply" in spin_r2:
                params["spin_r2_apply"] = spin_r2["apply"]
            if "learn" in spin_r2:
                params["spin_r2_learn"] = spin_r2["learn"]
            if "matrixPath" in spin_r2:
                params["spin_r2_matrix_path"] = spin_r2["matrixPath"]
 
        # QKRotation
        if "QKRotation" in config_dict:
            qk_rotation = config_dict["QKRotation"]
            if "apply" in qk_rotation:
                params["qk_rotation_apply"] = qk_rotation["apply"]
            if "matrixPath" in qk_rotation:
                params["qk_rotation_matrix_path"] = qk_rotation["matrixPath"]
 
        # FlattenQuant
        if "FlattenQuant" in config_dict:
            flatten_quant = config_dict["FlattenQuant"]
            if "apply" in flatten_quant:
                params["flatten_quant_apply"] = flatten_quant["apply"]
            if "learn" in flatten_quant:
                params["flatten_quant_learn"] = flatten_quant["learn"]
            if "applyThreshold" in flatten_quant:
                params["flatten_quant_apply_threshold"] = flatten_quant[
                    "applyThreshold"
                ]
            if "maxOverhead" in flatten_quant:
                params["flatten_quant_max_overhead"] = flatten_quant["maxOverhead"]
 
        # OptimizeFFN
        if "OptimizeFFN" in config_dict:
            optimize_ffn = config_dict["OptimizeFFN"]
            if "apply" in optimize_ffn:
                params["optimize_ffn_apply"] = optimize_ffn["apply"]
            if "chPerFFN" in optimize_ffn:
                params["optimize_ffn_ch_per_ffn"] = optimize_ffn["chPerFFN"]
 
        return cls(**params)
 
    def to_dict(self):
        return {
            "seed": self.seed,
            "applyHadamardRotationMatrix": self.apply_hadamard_rotation_matrix,
            "NormConv": {
                "apply": self.norm_conv_apply,
                "learn": self.norm_conv_learn,
                "smoothingFactor": self.norm_conv_smoothing_factor,
                "minGamma": self.norm_conv_min_gamma,
                "maxGamma": self.norm_conv_max_gamma,
            },
            "QK": {
                "apply": self.qk_apply,
                "smoothingFactor": self.qk_smoothing_factor,
                "minGamma": self.qk_min_gamma,
                "maxGamma": self.qk_max_gamma,
            },
            "UD": {
                "apply": self.ud_apply,
                "learn": self.ud_learn,
                "smoothingFactor": self.ud_smoothing_factor,
                "minGamma": self.ud_min_gamma,
                "maxGamma": self.ud_max_gamma,
            },
            "VO": {
                "apply": self.vo_apply,
                "smoothingFactor": self.vo_smoothing_factor,
                "minGamma": self.vo_min_gamma,
                "maxGamma": self.vo_max_gamma,
            },
            "FeedForwardMultiLUT": {
                "apply": self.ff_multi_lut_apply,
                "breakpoints": self.ff_multi_lut_breakpoints,
            },
            "SpinR1": {
                "apply": self.spin_r1_apply,
                "matrixPath": self.spin_r1_matrix_path,
            },
            "HeadOutChRotation": {
                "apply": self.head_out_ch_rotation_apply,
                "matrixPath": self.head_out_ch_rotation_matrix_path,
            },
            "SpinR2": {
                "apply": self.spin_r2_apply,
                "learn": self.spin_r2_learn,
                "matrixPath": self.spin_r2_matrix_path,
            },
            "QKRotation": {
                "apply": self.qk_rotation_apply,
                "matrixPath": self.qk_rotation_matrix_path,
            },
            "FlattenQuant": {
                "apply": self.flatten_quant_apply,
                "learn": self.flatten_quant_learn,
                "applyThreshold": self.flatten_quant_apply_threshold,
                "maxOverhead": self.flatten_quant_max_overhead,
            },
            "OptimizeFFN": {
                "apply": self.optimize_ffn_apply,
                "chPerFFN": self.optimize_ffn_ch_per_ffn,
            },
        }
 
 
def get_equivalent_transformation_config(**kwargs) -> EquivalentTransformationConfig:
    """
    @brief Create EquivalentTransformationConfig with partial parameters merged with defaults.
    """
    params = EquivalentTransformationConfig.DEFAULTS.copy()
    params.update(kwargs)
    return EquivalentTransformationConfig(
        **params, optional=EquivalentTransformationConfig.optional
    )
 
 
# }@