first commit

dev/functions/config_builder/__init__.py

@@ -0,0 +1,151 @@
+import torch
+
+from .qwen3moe import Qwen3MoEQuantConfigBuilder
+from .default import DefaultQuantConfigBuilder
+from .kimik2 import KimiK2QuantConfigBuilder
+from .glm4_7 import GLM4_7QuantConfigBuilder
+from ..specs import QuantArtifacts, ResolvedQuantConfig, DequantSpec
+
+
+def create_quant_config_builder(qconfig, custom_dequant_config, model_type):
+
+    if model_type in ["qwen3_moe", "minimax_m2", "minimax", "qwen3_5_moe", "deepseek_v3", "glm_moe_dsa"]:
+        return Qwen3MoEQuantConfigBuilder(qconfig, custom_dequant_config)
+    elif model_type == "kimik2":
+        return KimiK2QuantConfigBuilder(qconfig, custom_dequant_config)
+    elif model_type == "glm4_moe":
+        return GLM4_7QuantConfigBuilder(qconfig, custom_dequant_config)
+
+    return DefaultQuantConfigBuilder(
+        qconfig,
+        custom_dequant_config
+    )
+
+
+def derive_quant_key(weight_name: str, key_type: str, tensors: dict[str, torch.Tensor]) -> str | None:
+    """
+    根据 weight_name 推导量化辅助 tensor 名称。
+    """
+    # 定义所有可能的原始 weight 后缀
+    ORIG_WEIGHT_SUFFIXES = QWEIGHT_SUFFIXES
+
+    # 选择目标后缀表
+    if key_type == "scale":
+        target_suffixes = SCALE_SUFFIXES
+    elif key_type == "zero_point":
+        target_suffixes = ZERO_POINT_SUFFIXES
+    elif key_type == "qweight":
+        target_suffixes = QWEIGHT_SUFFIXES
+    elif key_type == "shape":
+        target_suffixes = SHAPE_SUFFIXES
+    else:
+        raise ValueError(f"Unknown key_type {key_type}")
+
+    candidates = []
+    for orig_suffix in ORIG_WEIGHT_SUFFIXES:
+        if weight_name.endswith(orig_suffix):
+            prefix = weight_name[: -len(orig_suffix)]
+            candidates.extend([prefix + s for s in target_suffixes])
+            break
+    else:
+        # weight_name 不匹配任何已知 weight 后缀，直接加 target_suffix
+        candidates.extend([weight_name + s for s in target_suffixes])
+
+    # 返回第一个存在于 tensors 的 key
+    for k in candidates:
+        if k in tensors:
+            return k
+    return None
+
+
+SCALE_SUFFIXES = [
+    ".weight_scale",
+    ".weight_scale_inv",
+    ".weight_block_scale",
+    "_scale",
+    "_scale_inv",
+    "_block_scale",
+]
+
+ZERO_POINT_SUFFIXES = [
+    ".weight_zero_point",
+    ".weight_zp",
+    "_zero_point",
+    "_zp",
+]
+
+QWEIGHT_SUFFIXES = [
+    ".weight",
+    ".weight_packed",
+    ".weight_q",
+    "_q",
+    "_packed",
+]
+
+SHAPE_SUFFIXES = [
+    ".weight_shape",
+    ".weight_block_shape",
+    "_shape",
+    "_block_shape",
+]
+
+
+def resolve_quant_artifacts(
+    weight_name: str,
+    tensors: dict[str, torch.Tensor],
+    rq: ResolvedQuantConfig,
+) -> QuantArtifacts:
+
+    if weight_name not in tensors:
+        raise KeyError(f"Missing quantized weight: {weight_name}")
+
+    if any(
+        weight_name.endswith(suffix)
+        for suffix in SCALE_SUFFIXES + ZERO_POINT_SUFFIXES + SHAPE_SUFFIXES
+    ):
+        return None
+
+    qweight_key = derive_quant_key(weight_name, "qweight", tensors)
+    scale_key = derive_quant_key(weight_name, "scale", tensors)
+    zero_point_key = derive_quant_key(weight_name, "zero_point", tensors)
+    shape_key = derive_quant_key(weight_name, "shape", tensors)
+
+    try:
+        qweight = tensors[qweight_key]
+    except KeyError:
+        raise KeyError(f"Missing quantized weight: {weight_name}")
+
+    # if qweight_key is None:
+    #     raise KeyError(f"Cannot find qweight for {weight_name}")
+
+    qweight = tensors[qweight_key]
+
+    if not scale_key:
+        raise KeyError(f"Missing scale tensor: {weight_name}")
+
+    scale = tensors[scale_key]
+
+    zero_point = None
+    if not rq.symmetric:
+        if zero_point_key is None:
+            raise KeyError(f"Cannot find zero_point tensor for {weight_name}")
+        zero_point = tensors[zero_point_key]
+
+    # shape = tensors.get(shape_key, None)
+    shape = None
+    if shape_key is not None:
+        shape = tensors[shape_key]
+
+    return QuantArtifacts(
+        qweight=qweight, scale=scale, zero_point=zero_point, shape=shape
+    )
+
+
+def to_dequant_spec(rq: ResolvedQuantConfig) -> DequantSpec:
+    return DequantSpec(
+        dtype=rq.dtype,
+        num_bits=rq.num_bits,
+        packed=rq.packed,
+        strategy=rq.strategy,
+        symmetric=rq.symmetric,
+    )

dev/functions/config_builder/base.py

@@ -0,0 +1,36 @@
+# config_builder/base.py
+import re
+
+from typing import Optional
+from ..specs import ResolvedQuantConfig
+
+
+class QuantConfigBuilder:
+
+    def resolve(
+        self,
+        weight_name: str,
+    ) -> Optional[ResolvedQuantConfig]:
+        pass
+    
+    def get_ignore(self, quantization_config: dict):
+        ignore = []
+
+        candidate_suffixes = ["modules_to_not_convert", "ignore"]
+        for cs in candidate_suffixes:
+            ig = quantization_config.get(cs)
+            if ig:
+                ignore.extend(ig)
+
+        return ignore
+    
+    def compile_ignore(self, ignore_list):
+        exact = set()
+        regex = []
+
+        for x in (ignore_list or []):
+            if x.startswith("re:"):
+                regex.append(re.compile(x[3:]))
+            else:
+                exact.add(x)
+        return list(exact), regex

dev/functions/config_builder/default.py

@@ -0,0 +1,56 @@
+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)
+        )

dev/functions/config_builder/glm4_7.py

@@ -0,0 +1,39 @@
+# config_builder/qwen3.py
+import re
+
+from .base import QuantConfigBuilder
+from ..specs import ResolvedQuantConfig
+
+
+class GLM4_7QuantConfigBuilder(QuantConfigBuilder):
+
+    def __init__(self, qconfig: dict, custom_dequant_config: dict):
+        qc = qconfig
+        # get ignore from qconfig
+        self.original_ignore = self.get_ignore(qc)
+        # get ignore from custom_dequant_config
+        self.custom_ignore = self.get_ignore(custom_dequant_config)
+        self.exact_ignore, self.regex_ignore = self.compile_ignore(self.original_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 resolve(self, weight_name: str):
+        if self._ignored(weight_name):
+            return None
+
+        return ResolvedQuantConfig(
+            dtype="float",
+            num_bits=8,
+            packed=False,
+            strategy="channel",
+            group_size=None,
+            block_structure=(0,0),
+            symmetric=True,
+            dynamic=False,
+        )

dev/functions/config_builder/kimik2.py

@@ -0,0 +1,38 @@
+# config_builder/qwen3.py
+import re
+
+from .base import QuantConfigBuilder
+from ..specs import ResolvedQuantConfig
+
+
+class KimiK2QuantConfigBuilder(QuantConfigBuilder):
+
+    def __init__(self, qconfig: dict, custom_dequant_config: dict):
+        qc = qconfig
+        self.group_size = qconfig['config_groups']['group_0']['weights']['group_size']
+        self.original_ignore = self.get_ignore(qc)
+        self.custom_ignore = self.get_ignore(custom_dequant_config)
+        self.exact_ignore, self.regex_ignore = self.compile_ignore(self.original_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 resolve(self, weight_name: str):
+        if self._ignored(weight_name):
+            return None
+
+        return ResolvedQuantConfig(
+            dtype="int",
+            num_bits=4,
+            packed=True,
+            strategy="group",
+            group_size=self.group_size,
+            block_structure=None,
+            symmetric=True,
+            dynamic=False,
+        )

dev/functions/config_builder/kimik25.py

@@ -0,0 +1,38 @@
+# config_builder/qwen3.py
+import re
+
+from .base import QuantConfigBuilder
+from ..specs import ResolvedQuantConfig
+
+
+class KimiK25QuantConfigBuilder(QuantConfigBuilder):
+
+    def __init__(self, qconfig: dict, custom_dequant_config: dict):
+        qc = qconfig
+        self.group_size = qconfig['config_groups']['group_0']['weights']['group_size']
+        self.original_ignore = self.get_ignore(qc)
+        self.custom_ignore = self.get_ignore(custom_dequant_config)
+        self.exact_ignore, self.regex_ignore = self.compile_ignore(self.original_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 resolve(self, weight_name: str):
+        if self._ignored(weight_name):
+            return None
+
+        return ResolvedQuantConfig(
+            dtype="int",
+            num_bits=4,
+            packed=True,
+            strategy="group",
+            group_size=self.group_size,
+            block_structure=None,
+            symmetric=True,
+            dynamic=False,
+        )

dev/functions/config_builder/qwen3moe.py

@@ -0,0 +1,40 @@
+# config_builder/qwen3.py
+import re
+
+from .base import QuantConfigBuilder
+from ..specs import ResolvedQuantConfig
+
+
+class Qwen3MoEQuantConfigBuilder(QuantConfigBuilder):
+
+    def __init__(self, qconfig: dict, custom_dequant_config: dict):
+        qc = qconfig
+        self.block_h, self.block_w = qc["weight_block_size"]
+        # get ignore from qconfig
+        self.original_ignore = self.get_ignore(qc)
+        # get ignore from custom_dequant_config
+        self.custom_ignore = self.get_ignore(custom_dequant_config)
+        self.exact_ignore, self.regex_ignore = self.compile_ignore(self.original_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 resolve(self, weight_name: str):
+        if self._ignored(weight_name):
+            return None
+
+        return ResolvedQuantConfig(
+            dtype="float",
+            num_bits=8,
+            packed=False,
+            strategy="block",
+            group_size=None,
+            block_structure=(self.block_h, self.block_w),
+            symmetric=True,
+            dynamic=False,
+        )