57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
|
import re
|
||
|
|
from .base import QuantConfigBuilder
|
||
|
|
from ..specs import ResolvedQuantConfig
|
||
|
|
|
||
|
|
|
||
|
|
class DefaultQuantConfigBuilder(QuantConfigBuilder):
|
||
|
|
|
||
|
|
def __init__(self, qconfig: dict, custom_dequant_config: dict):
|
||
|
|
self.qconfig = qconfig
|
||
|
|
self.groups = qconfig["config_groups"]
|
||
|
|
self.custom_ignore = self.get_ignore(custom_dequant_config)
|
||
|
|
self.exact_ignore, self.regex_ignore = self.compile_ignore(self.custom_ignore)
|
||
|
|
|
||
|
|
def _ignored(self, weight_name: str) -> bool:
|
||
|
|
if any(weight_name.startswith(ig) for ig in self.exact_ignore):
|
||
|
|
return True
|
||
|
|
if any(re.match(p, weight_name) for p in self.regex_ignore):
|
||
|
|
return True
|
||
|
|
|
||
|
|
return False
|
||
|
|
|
||
|
|
def _match_group(self, weight_name: str):
|
||
|
|
for g in self.groups.values():
|
||
|
|
for p in g.get("weight_patterns", []):
|
||
|
|
if p.startswith("re:") and re.match(p[3:], weight_name):
|
||
|
|
return g
|
||
|
|
return None
|
||
|
|
|
||
|
|
def build_weight_quant_config(self, qconfig: dict, group_cfg: dict) -> dict:
|
||
|
|
wcfg = group_cfg["weights"]
|
||
|
|
|
||
|
|
return {
|
||
|
|
"dtype": wcfg["type"], # int / float
|
||
|
|
"num_bits": wcfg["num_bits"],
|
||
|
|
"packed": qconfig["format"] == "pack-quantized",
|
||
|
|
"strategy": wcfg["strategy"],
|
||
|
|
"group_size": wcfg.get("group_size"),
|
||
|
|
"block_structure": wcfg.get("block_structure"),
|
||
|
|
"symmetric": wcfg["symmetric"],
|
||
|
|
"dynamic": wcfg["dynamic"],
|
||
|
|
}
|
||
|
|
|
||
|
|
def resolve(self, weight_name: str):
|
||
|
|
if self._ignored(weight_name):
|
||
|
|
return None
|
||
|
|
|
||
|
|
if len(self.groups) == 1:
|
||
|
|
group = next(iter(self.groups.values()))
|
||
|
|
else:
|
||
|
|
group = self._match_group(weight_name)
|
||
|
|
if group is None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
return ResolvedQuantConfig(
|
||
|
|
**self.build_weight_quant_config(self.qconfig, group)
|
||
|
|
)
|