saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add int4 quantization code in python (#17077)

Browse source 
### Description
Adding int4 quantization code in python


### Motivation and Context
Python quantization tool no-longer needs to invoke shell to call a
native exe
This commit is contained in:
Chen Fu 2023-08-11 15:17:58 -07:00 committed by GitHub
parent 5704e71b89
commit f2e1b91634
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 244 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
onnxruntime/python/tools/quantization/__init__.py
View file

@ -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

99
onnxruntime/python/tools/quantization/matmul_weight4_quantizer.py
View file

@ -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)

45
onnxruntime/python/tools/quantization/q4dq_wrapper.py
View file

@ -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

153
onnxruntime/test/python/quantization/test_op_matmulfpq4.py Normal file
View file

@ -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()