diff --git a/onnxruntime/python/tools/quantization/__init__.py b/onnxruntime/python/tools/quantization/__init__.py index 170c0928fe..9d397499d4 100644 --- a/onnxruntime/python/tools/quantization/__init__.py +++ b/onnxruntime/python/tools/quantization/__init__.py @@ -5,7 +5,6 @@ from .calibrate import ( # noqa: F401 MinMaxCalibrater, create_calibrator, ) -from .matmul_weight4_quantizer import MatMulWeight4Quantizer # noqa: F401 from .qdq_quantizer import QDQQuantizer # noqa: F401 from .quant_utils import QuantFormat, QuantType, write_calibration_table # noqa: F401 from .quantize import DynamicQuantConfig # noqa: F401 diff --git a/onnxruntime/python/tools/quantization/matmul_weight4_quantizer.py b/onnxruntime/python/tools/quantization/matmul_weight4_quantizer.py deleted file mode 100644 index 921e02fb69..0000000000 --- a/onnxruntime/python/tools/quantization/matmul_weight4_quantizer.py +++ /dev/null @@ -1,260 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. See License.txt in the project root for -# license information. -# -------------------------------------------------------------------------- - -import argparse -import struct -from pathlib import Path -from typing import List, Tuple - -import numpy as np -import numpy.typing as npt -import onnx -from onnx.onnx_pb import GraphProto, ModelProto, NodeProto, TensorProto - -from .onnx_model import ONNXModel -from .quant_utils import attribute_to_kwarg, load_model_with_shape_infer - - -def __q4_block_size(quant_type: int) -> int: - # happens to be 32 for now, but future quantization types - # may have bigger block size - return 32 - - -def __q4_blob_size(quant_type: int) -> int: - if quant_type == MatMulWeight4Quantizer.BlkQ4Sym: - # 4b each value, with one fp32 scale - blob_size = 32 // 2 + 4 - elif quant_type == MatMulWeight4Quantizer.BlkQ4Zp8: - # 4b each value, with one fp32 scale and one uint8 zero point - blob_size = 32 // 2 + 4 + 1 - else: - raise ValueError(f"Unsupported quantization type: {quant_type}") - return blob_size - - -def __q4_buf_size(quant_type: int, rows: int, cols: int) -> int: - block_size = __q4_block_size(quant_type) - blob_size = __q4_blob_size(quant_type) - k_blocks = (rows + block_size - 1) // block_size - return k_blocks * cols * blob_size - - -def int4_block_quant(quant_type: int, fp32weight: npt.ArrayLike) -> np.ndarray: - """4b quantize fp32 weight to a blob""" - - if len(fp32weight.shape) != 2: - raise ValueError("Current int4 block quantization only supports 2D tensors!") - rows, cols = fp32weight.shape - - block_size = __q4_block_size(quant_type) - blob_size = __q4_blob_size(quant_type) - k_blocks = (rows + block_size - 1) // block_size - padded_rows = k_blocks * block_size - pad_len = padded_rows - rows - if pad_len > 0: - fp32weight = np.pad(fp32weight, ((0, pad_len), (0, 0)), "constant") - - # block wise quantization, each block comes from a single column - blob_idx = 0 - packed = np.zeros((cols * k_blocks, blob_size), dtype="uint8") - for n in range(cols): - ncol = fp32weight[:, n] - blks = np.split(ncol, k_blocks) - for blk in blks: - packed_blob = packed[blob_idx] - blob_idx += 1 - - if quant_type == MatMulWeight4Quantizer.BlkQ4Sym: - amax_idx = np.argmax(np.abs(blk)) - bmax = blk[amax_idx] - scale = bmax / (-8) - zp = 8 - else: - vmin = np.min(blk) - vmax = np.max(blk) - vmin = min(vmin, 0.0) - vmax = max(vmax, 0.0) - scale = (vmax - vmin) / ((1 << 4) - 1) - zero_point_fp = vmin - if scale != 0.0: - zero_point_fp = 0.0 - vmin / scale - zp = min(15, max(0, round(zero_point_fp))) - - reciprocal_scale = 1.0 / scale if scale != 0 else 0.0 - bf = struct.pack("f", scale) - packed_blob[0] = bf[0] - packed_blob[1] = bf[1] - packed_blob[2] = bf[2] - packed_blob[3] = bf[3] - blob_offset = 4 - if quant_type == MatMulWeight4Quantizer.BlkQ4Zp8: - packed_blob[4] = zp - blob_offset = 5 - - num_segs = block_size // 32 - blk_int = np.clip(np.rint(blk * reciprocal_scale + zp), 0, 15).astype("uint8") - segs = np.split(blk_int, num_segs) - for seg in segs: - packed_blob[blob_offset : (blob_offset + 16)] = np.bitwise_or(seg[0:16], np.left_shift(seg[16:32], 4)) - blob_offset += 16 - return packed.reshape(-1) - - -class MatMulWeight4Quantizer: - """Perform 4b quantization of constant MatMul weights""" - - ################## - # quantization types, must be consistent with native code type - # MLAS_BLK_QUANT_TYPE defined in mlas_q4.h - - # 32 number block, symmetric quantization, with one fp32 as scale, zero point is always 0 - BlkQ4Sym = 0 - - # 32 number block, quantization, with one fp32 as scale, one uint8 zero point - BlkQ4Zp8 = 1 - - def __init__(self, model: ModelProto, quant_type: int): - self.model = ONNXModel(model) - self.quant_type = quant_type - - @staticmethod - def __get_initializer(name, graph_path: List[GraphProto]) -> Tuple[TensorProto, GraphProto]: - for gid in range(len(graph_path) - 1, -1, -1): - graph = graph_path[gid] - for tensor in graph.initializer: - if tensor.name == name: - return tensor, graph - return None, None - - def _q4_matmul_node_weight(self, node: NodeProto, graph_stack: List[GraphProto]) -> NodeProto: - """If the node is MatMul with fp32 const weight, quantize the weight with int4, and return the new node""" - - if node.op_type != "MatMul": - return node # only care about MatMul for now - - inputB = node.input[1] # noqa: N806 - B, Bs_graph = MatMulWeight4Quantizer.__get_initializer(inputB, graph_stack) # noqa: N806 - if B is None: - return node # only care about constant weight - - # TODO!! assume B is not used by any other node - B_array = onnx.numpy_helper.to_array(B) # noqa: N806 - if len(B_array.shape) != 2: - return node # can only process 2-D matrix - - rows, cols = B_array.shape - packed = int4_block_quant(self.quant_type, B_array) - B_quant = onnx.numpy_helper.from_array(packed) # noqa: N806 - B_quant.name = B.name + "_Q4" - Bs_graph.initializer.remove(B) - for input in Bs_graph.input: - if input.name == inputB: - Bs_graph.input.remove(input) - break - - B_shape = onnx.numpy_helper.from_array(np.array([rows, cols]).astype(np.int64)) # noqa: N806 - B_shape.name = B.name + "_shape" - Bs_graph.initializer.extend([B_quant, B_shape]) - - kwargs = {} - kwargs["blk_quant_type"] = self.quant_type - matmul_q4_node = onnx.helper.make_node( - "MatMulFpQ4", - inputs=[node.input[0], B_quant.name, B_shape.name], - outputs=[node.output[0]], - name=node.name + "_Q4" if node.name else "", - domain="com.microsoft", - **kwargs, - ) - return matmul_q4_node - - def _process_subgraph(self, graph_stack: List[GraphProto]): - new_nodes = [] - graph = graph_stack[-1] - - for node in graph.node: - graph_attrs = [ - attr - for attr in node.attribute - if attr.type == onnx.AttributeProto.GRAPH or attr.type == onnx.AttributeProto.GRAPHS - ] - if len(graph_attrs): - kwargs = {} - for attr in node.attribute: - if attr.type == onnx.AttributeProto.GRAPH: - # recursive call to take care of sub-graph - graph_stack.append(attr.g) - kv = {attr.name: self._process_subgraph(graph_stack)} - elif attr.type == onnx.AttributeProto.GRAPHS: - value = [] - for subgraph in attr.graphs: - # recursive call to take care of sub-graph - graph_stack.append(subgraph) - value.extend([self._process_subgraph(graph_stack)]) - kv = {attr.name: value} - else: - kv = attribute_to_kwarg(attr) - kwargs.update(kv) - node = onnx.helper.make_node( # noqa: PLW2901 - node.op_type, node.input, node.output, name=node.name, **kwargs - ) - - new_nodes.append(self._q4_matmul_node_weight(node, graph_stack)) - - graph.ClearField("node") - graph.node.extend(new_nodes) - graph_stack.pop() - return graph - - def process(self): - # use a stack to keep track of sub-graphs - graph_stack = [self.model.graph()] - opset_import = self.model.opset_import() - - has_ms_domain = False - for opset in opset_import: - if opset.domain == "com.microsoft": - has_ms_domain = True - if not has_ms_domain: - opset_import.extend([onnx.helper.make_opsetid("com.microsoft", 1)]) - - self._process_subgraph(graph_stack) - - -def parse_args(): - parser = argparse.ArgumentParser( - description="""Blockwise int4 quantization for MatMul 2D weight matrices. - -A weight matrix is partitioned into into blocks, where each block is a -continguous subset inside each column. Each block is quantized into a -set of 4b integers with a scaling factor and an optional offset. -""" - ) - - parser.add_argument("--input_model", required=True, help="Path to the input model file") - parser.add_argument("--output_model", required=True, help="Path to the output model file") - parser.add_argument( - "--quant_bin_path", - required=True, - help="""Currently quantization code is implemented in a separate binary -(onnxruntime_mlas_q4dq) that is compiled with Onnxruntime native code. -Path to this binary needs to be provided here.""", - ) - return parser.parse_args() - - -if __name__ == "__main__": - args = parse_args() - - input_model_path = args.input_model - output_model_path = args.output_model - q4dq_bin_path = args.quant_bin_path - - model = load_model_with_shape_infer(Path(input_model_path)) - quant = MatMulWeight4Quantizer(model, 0) - quant.process() - quant.model.save_model_to_file(output_model_path, False) diff --git a/onnxruntime/test/python/quantization/test_op_matmulfpq4.py b/onnxruntime/test/python/quantization/test_op_matmulfpq4.py deleted file mode 100644 index 170bb09a0f..0000000000 --- a/onnxruntime/test/python/quantization/test_op_matmulfpq4.py +++ /dev/null @@ -1,153 +0,0 @@ -#!/usr/bin/env python -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. See License.txt in the project root for -# license information. -# -------------------------------------------------------------------------- - -import tempfile -import unittest -from pathlib import Path -from typing import Dict, Tuple, Union - -import numpy as np -import onnx -from onnx import TensorProto, helper -from op_test_utils import TestDataFeeds, check_model_correctness, check_op_type_count - -from onnxruntime.quantization import MatMulWeight4Quantizer, quant_utils - - -class TestOpMatMulFpQ4(unittest.TestCase): - @classmethod - def setUpClass(cls): - cls._tmp_model_dir = tempfile.TemporaryDirectory(prefix="test_matmulfpq4.") - - @classmethod - def tearDownClass(cls): - cls._tmp_model_dir.cleanup() - - def fill_int4_data(self, shape: Union[int, Tuple[int, ...]], symmetric: bool) -> np.ndarray: - line = np.zeros(shape) - line = line.reshape(-1) - - if symmetric: - v = -2.0 - for i in range(line.shape[0]): - if v == 0 or v == -3 or v == 3: - v += 1 - line[i] = v - v += 1 - if v >= 8: - v = -8 - else: - v = 0.0 - for i in range(line.shape[0]): - line[i] = v - v += 1 - if v >= 16: - v = 0 - - return line.reshape(shape) - - def input_feeds(self, n: int, name2shape: Dict[str, Union[int, Tuple[int, ...]]]) -> TestDataFeeds: - input_data_list = [] - for _i in range(n): - inputs = {} - for name, shape in name2shape.items(): - inputs.update({name: np.random.randint(-1, 2, shape).astype(np.float32)}) - input_data_list.extend([inputs]) - dr = TestDataFeeds(input_data_list) - return dr - - def construct_model_matmul(self, output_model_path: str, symmetric: bool) -> None: - # (input) - # | - # MatMul - # | - # (output) - input_name = "input" - output_name = "output" - initializers = [] - - def make_gemm(input_name, weight_shape: Union[int, Tuple[int, ...]], weight_name: str, output_name: str): - weight_data = self.fill_int4_data(weight_shape, symmetric).astype(np.float32) - initializers.append(onnx.numpy_helper.from_array(weight_data, name=weight_name)) - return onnx.helper.make_node( - "MatMul", - [input_name, weight_name], - [output_name], - ) - - in_features = 52 - out_features = 288 - # make MatMulFpQ4 node - matmul_node = make_gemm( - input_name, - [in_features, out_features], - "linear1.weight", - output_name, - ) - - # make graph - input_tensor = helper.make_tensor_value_info(input_name, TensorProto.FLOAT, [-1, in_features]) - output_tensor = helper.make_tensor_value_info(output_name, TensorProto.FLOAT, [-1, out_features]) - graph_name = "matmul_test" - graph = helper.make_graph( - [matmul_node], - graph_name, - [input_tensor], - [output_tensor], - initializer=initializers, - ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) - model.ir_version = 7 # use stable onnx ir version - - onnx.save(model, output_model_path) - - def quant_test( - self, - model_fp32_path: str, - data_reader: TestDataFeeds, - quantization_type: int, # 0: BlkQ4Sym, 1: BlkQ4Zp8 - ): - qtype_str = "BlkQ4Sym" if (quantization_type == 0) else "BlkQ4Zp8" - model_int4_path = str(Path(self._tmp_model_dir.name).joinpath(f"matmulfpq4_{qtype_str}.onnx").absolute()) - - # Quantize fp32 model to int4 model - model = quant_utils.load_model_with_shape_infer(Path(model_fp32_path)) - quant = MatMulWeight4Quantizer(model, quantization_type) - quant.process() - quant.model.save_model_to_file(model_int4_path, False) - - quant_nodes = {"MatMulFpQ4": 1} - check_op_type_count(self, model_int4_path, **quant_nodes) - - data_reader.rewind() - - try: - check_model_correctness(self, model_fp32_path, model_int4_path, data_reader.get_next()) - except Exception as exception: - if "4b quantization not yet supported on this hardware platform!" in exception.args[0]: - # Currently we don't have int4 quantization support on all platforms, has to tolerate this exception - pass - else: - raise exception - - def test_quantize_matmul_int4_symmetric(self): - np.random.seed(13) - - model_fp32_path = str(Path(self._tmp_model_dir.name).joinpath("matmul_fp32_symmetric.onnx").absolute()) - self.construct_model_matmul(model_fp32_path, symmetric=True) - data_reader = self.input_feeds(1, {"input": [100, 52]}) - self.quant_test(model_fp32_path, data_reader, quantization_type=MatMulWeight4Quantizer.BlkQ4Sym) - - def test_quantize_matmul_int4_offsets(self): - model_fp32_path = str(Path(self._tmp_model_dir.name).joinpath("matmul_fp32_offset.onnx").absolute()) - self.construct_model_matmul(model_fp32_path, symmetric=False) - data_reader = self.input_feeds(1, {"input": [100, 52]}) - self.quant_test(model_fp32_path, data_reader, quantization_type=MatMulWeight4Quantizer.BlkQ4Zp8) - - -if __name__ == "__main__": - unittest.main()