diff --git a/onnxruntime/python/tools/quantization/__init__.py b/onnxruntime/python/tools/quantization/__init__.py index bd9494ccbc..170c0928fe 100644 --- a/onnxruntime/python/tools/quantization/__init__.py +++ b/onnxruntime/python/tools/quantization/__init__.py @@ -6,7 +6,6 @@ from .calibrate import ( # noqa: F401 create_calibrator, ) from .matmul_weight4_quantizer import MatMulWeight4Quantizer # noqa: F401 -from .q4dq_wrapper import Q4dqWrapper # 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 index 9f1158c4c6..44d870bb22 100644 --- a/onnxruntime/python/tools/quantization/matmul_weight4_quantizer.py +++ b/onnxruntime/python/tools/quantization/matmul_weight4_quantizer.py @@ -5,18 +5,105 @@ # -------------------------------------------------------------------------- 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 .q4dq_wrapper import Q4dqWrapper 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""" @@ -30,9 +117,8 @@ class MatMulWeight4Quantizer: # 32 number block, quantization, with one fp32 as scale, one uint8 zero point BlkQ4Zp8 = 1 - def __init__(self, model: ModelProto, q4dq: Q4dqWrapper, quant_type: int): + def __init__(self, model: ModelProto, quant_type: int): self.model = ONNXModel(model) - self.q4dq = q4dq self.quant_type = quant_type @staticmethod @@ -61,8 +147,7 @@ class MatMulWeight4Quantizer: return node # can only process 2-D matrix rows, cols = B_array.shape - packed = self.q4dq.quantize(B_array, self.quant_type) - + 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) @@ -169,9 +254,7 @@ if __name__ == "__main__": output_model_path = args.output_model q4dq_bin_path = args.quant_bin_path - q4dq = Q4dqWrapper(q4dq_bin_path) - model = load_model_with_shape_infer(Path(input_model_path)) - quant = MatMulWeight4Quantizer(model, q4dq, 0) + quant = MatMulWeight4Quantizer(model, 0) quant.process() quant.model.save_model_to_file(output_model_path, False) diff --git a/onnxruntime/python/tools/quantization/q4dq_wrapper.py b/onnxruntime/python/tools/quantization/q4dq_wrapper.py deleted file mode 100644 index 62cc6d0ea3..0000000000 --- a/onnxruntime/python/tools/quantization/q4dq_wrapper.py +++ /dev/null @@ -1,45 +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 os -import subprocess -import tempfile - -import numpy as np -import numpy.typing as npt - - -class Q4dqWrapper: - """A wrapper to native command line onnxruntime_mlas_q4dq""" - - def __init__(self, exepath: str): - self.q4dq_cmd = exepath - - def quantize(self, fp32weight: npt.ArrayLike, quant_type: int) -> np.ndarray: - """4b quantize fp32 weight to a blob""" - - array = fp32weight.astype(np.float32) - if len(array.shape) != 2: - raise Exception("Only 2D fp32 array accepted!") - rows, cols = array.shape - - with tempfile.TemporaryDirectory() as tmpdirname: - fp32file = os.path.join(tmpdirname, "fp32weight") - array.tofile(fp32file) - - q4file = os.path.join(tmpdirname, "q4weight") - - cmd = "{cmdpath} q {k} {n} --quant_type {qtype} --input_file {fp32} --output_file {q4} --output_format bin".format( - cmdpath=self.q4dq_cmd, k=rows, n=cols, qtype=quant_type, fp32=fp32file, q4=q4file - ) - subprocess.run(cmd, shell=True) - - if not os.path.isfile(q4file): - raise Exception("Quantization failed, 4b quantization is not yet supported on this platform!") - - packed = np.fromfile(q4file, dtype="uint8") - return packed diff --git a/onnxruntime/test/python/quantization/test_op_matmulfpq4.py b/onnxruntime/test/python/quantization/test_op_matmulfpq4.py new file mode 100644 index 0000000000..170bb09a0f --- /dev/null +++ b/onnxruntime/test/python/quantization/test_op_matmulfpq4.py @@ -0,0 +1,153 @@ +#!/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()