Method stream_quantize

pymllm/quantize/solver.py:77–174

(
        self, quant_cfg, tensor_dict: Dict, writer: ModelFileV2, **kwargs
    )

        return len(aux)

    def stream_quantize(
        self, quant_cfg, tensor_dict: Dict, writer: ModelFileV2, **kwargs
    ) -> bool:
        if not isinstance(writer, ModelFileV2):
            raise NotImplementedError(
                "stream_quantize only support type: ModelFileV2 currently."
            )

        if quant_cfg is None:
            quant_cfg = {}

        # Planning
        param_groups: Dict[str, List[Any, Dict]] = {}
        for k, v in quant_cfg.items():
            sub_group: Dict[str, QuantizePlanPayload] = {}
            hints = v["hints"]
            pattern = re.compile(k)
            for pk, pv in tensor_dict.items():
                if pattern.fullmatch(pk) is not None:
                    # pk is model.linear_0.weight or model.linear_0.bias
                    # layer_name is model.linear_0
                    layer_name, _ = pk.rsplit(".", 1)
                    if layer_name not in sub_group:
                        sub_group[layer_name] = QuantizePlanPayload()
                    sub_group[layer_name].inputs_num += 1
                    sub_group[layer_name].inputs_dict.update({pk: pv})
            param_groups.update({k: [hints, sub_group]})

        # Prepare inputs and outputs, calculate the params
        for group_regex in param_groups.keys():
            sub_group = param_groups[group_regex]
            for payload_name, payload in sub_group[1].items():
                for pass_ in self.passes:
                    if pass_.match(param_groups[group_regex][0], payload.inputs_dict):
                        prepared_payload = pass_.prepare(
                            param_groups[group_regex][0],
                            payload.inputs_dict,
                        )
                        param_groups[group_regex][1][payload_name] = prepared_payload
                        break

        # Show Planned Info
        verbose = kwargs.get("verbose", False)
        if verbose:
            print("Planned Quantized Info:")
            for group_regex in param_groups.keys():
                print(f"{group_regex}:")
                sub_group = param_groups[group_regex]
                for payload_name, payload in sub_group[1].items():
                    print(" " * 4 + payload_name + ":")
                    print(" " * 8 + f"inputs num: {payload.inputs_num}")
                    print(" " * 8 + f"outputs num: {payload.outputs_num}")
                    print(" " * 8 + "params before quantization:")
                    for k in payload.inputs_dict.keys():
                        print(" " * 12 + k)
                    print(" " * 8 + "params after quantization:")
                    for k in payload.outputs_dict.keys():
                        print(" " * 12 + k)