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

ZCode FastFormers changes (#5827)

Browse source 
* Add FBGEMM submodule

* Add fbgemm based per-channel quantization

* Add missing logic for pre-layernorm transformer model fusion

* add support for structured pruning architecture -fastformers

* Fix windows build

* Add a default behavior when head_size is not present for the backward compatibility

* Remove FBGEMM and default to tensor-wise quantization, column-wise quantization will be enabled later

* Fixed some unit test errors

* Fix windows compile error and unit test errors

* delete the option removed from the upstream

* Addresses review comments and fixes a merge error

* Remove commented out code

* add non-zero zp support

* support A and B scale with any dimensions

* fix build breaks

* fix warning in MSVC

* Fix bug for not checking original float value names when treat it as not existing.

* Clean up head size

* Clean up python tools

* Enable per column quantization

* fix quant weight cleanup bug

* A few code clean up

* Some code clean-up

* Some code clean-up

* Change option name

* update default value

* Rename option and parameter names

* Missing argument name change

* Add tests for quantization options for attention and matmul

Co-authored-by: Yufeng Li <liyufeng1987@gmail.com>
Co-authored-by: Lei Zhang <zhang.huanning@hotmail.com>
This commit is contained in:
Young Jin Kim 2021-05-17 21:12:21 -07:00 committed by GitHub
parent 38d90b0f15
commit e9057d2e49
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 112 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc
View file

@ -135,10 +135,10 @@ Status QAttention<T>::Compute(OpKernelContext* context) const {
// Input 1 - weights : (input_hidden_size, 3 * hidden_size)
// Input 2 - bias : (3 * hidden_size)
// Input 3 - input_scale : scalar
// Input 4 - weight_scale : scalar
// Input 4 - weight_scale : scalar for per tensor quantization, (3 * hidden_size) for per column quantization
// Input 5 - mask_index : nullptr, (batch_size), (2 * batch_size), (batch_size, 1), (1, 1) or (batch_size, past_sequence_length + sequence_length)
// Input 6 - input_zero_point : scalar
// Input 7 - weight_zero_point : scalar
// Input 7 - weight_zero_point : scalar for per tensor quantization, (3 * hidden_size) for per column quantization
// Input 8 - past : (2, batch_size, num_heads, past_sequence_length, head_size)
// Output 0 : (batch_size, sequence_length, hidden_size)
// Output 1 - present : (2, batch_size, num_heads, past_sequence_length + sequence_length, head_size)

34
onnxruntime/python/tools/quantization/onnx_quantizer.py
View file

@ -41,6 +41,8 @@ class ONNXQuantizer:
self.static = static # use static quantization for inputs.
self.fuse_dynamic_quant = False
self.extra_options = extra_options if extra_options is not None else {}
self.q_matmul_const_b_only = 'MatMulConstBOnly' in self.extra_options and self.extra_options['MatMulConstBOnly']
self.is_weight_symmetric = 'WeightSymmetric' not in self.extra_options or self.extra_options['WeightSymmetric']
self.input_qType = onnx_proto.TensorProto.INT8 if input_qType == QuantType.QInt8 else onnx_proto.TensorProto.UINT8
self.weight_qType = onnx_proto.TensorProto.INT8 if weight_qType == QuantType.QInt8 else onnx_proto.TensorProto.UINT8
@ -186,6 +188,11 @@ class ONNXQuantizer:
if self.nodes_to_exclude is not None and node.name in self.nodes_to_exclude:
return False
# do not quantize non-constant B matrices for matmul
if self.q_matmul_const_b_only:
if node.op_type == "MatMul" and find_by_name(node.input[1], self.model.initializer()) is None:
return False
return True
def quantize_model(self):
@ -590,7 +597,7 @@ class ONNXQuantizer:
return quantized_bias_name
def quantize_inputs(self, node, indices, initializer_use_weight_qType=True):
def quantize_inputs(self, node, indices, initializer_use_weight_qType=True, reduce_range=False, op_level_per_channel=False, axis=-1):
'''
Given a node, this function quantizes the inputs as follows:
- If input is an initializer, quantize the initializer data, replace old initializer
@ -623,8 +630,14 @@ class ONNXQuantizer:
# Quantize the input
initializer = find_by_name(node_input, self.model.initializer())
if initializer is not None:
q_weight_name, zp_name, scale_name = self.quantize_weight(
initializer, self.weight_qType if initializer_use_weight_qType else self.input_qType)
if self.per_channel and op_level_per_channel:
q_weight_name, zp_name, scale_name = self.quantize_weight_per_channel(
initializer.name, self.weight_qType if initializer_use_weight_qType else self.input_qType,
axis, reduce_range)
else:
q_weight_name, zp_name, scale_name = self.quantize_weight(
initializer, self.weight_qType if initializer_use_weight_qType else self.input_qType,
reduce_range)
quantized_input_names.append(q_weight_name)
zero_point_names.append(zp_name)
@ -649,7 +662,7 @@ class ONNXQuantizer:
return (quantized_input_names, zero_point_names, scale_names, nodes)
def quantize_weight(self, weight, qType):
def quantize_weight(self, weight, qType, reduce_range=False):
'''
:param weight: TensorProto initializer
:param qType: type to quantize to
@ -667,7 +680,8 @@ class ONNXQuantizer:
# Update packed weight, zero point, and scale initializers
weight_data = self.tensor_proto_to_array(weight)
_, _, zero_point, scale, q_weight_data = quantize_data(weight_data.flatten().tolist(),
get_qrange_for_qType(qType, self.reduce_range), qType)
get_qrange_for_qType(qType, self.reduce_range and reduce_range),
qType, self.is_weight_symmetric)
q_weight_data = np.asarray(q_weight_data, dtype=onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[qType]).reshape(weight.dims)
q_weight_initializer = onnx.numpy_helper.from_array(q_weight_data, q_weight_name)
@ -682,7 +696,7 @@ class ONNXQuantizer:
return q_weight_name, zp_name, scale_name
def quantize_weight_per_channel(self, weight_name, weight_qType, channel_axis):
def quantize_weight_per_channel(self, weight_name, weight_qType, channel_axis, reduce_range=True):
# Find if this input is already quantized
if weight_name in self.quantized_value_map:
quantized_value = self.quantized_value_map[weight_name]
@ -702,8 +716,9 @@ class ONNXQuantizer:
for i in range(channel_count):
per_channel_data = weights.take(i, channel_axis)
rmin, rmax, zero_point, scale, quantized_per_channel_data = quantize_data(
per_channel_data.flatten().tolist(), get_qrange_for_qType(weight_qType, self.reduce_range),
weight_qType)
per_channel_data.flatten().tolist(),
get_qrange_for_qType(weight_qType, self.reduce_range and reduce_range),
weight_qType, self.is_weight_symmetric)
rmin_list.append(rmin)
rmax_list.append(rmax)
zero_point_list.append(zero_point)
@ -802,7 +817,8 @@ class ONNXQuantizer:
rmax = max(rmax, 0)
quantization_params[tensor_name] = compute_scale_zp(rmin, rmax, self.input_qType,
get_qrange_for_qType(self.input_qType))
get_qrange_for_qType(self.input_qType),
self.is_weight_symmetric)
return quantization_params

2
onnxruntime/python/tools/quantization/operators/attention.py
View file

@ -21,7 +21,7 @@ class AttentionQuant(QuantOperatorBase):
assert (node.op_type == "Attention")
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer.quantize_inputs(node, [0, 1])
self.quantizer.quantize_inputs(node, [0, 1], reduce_range=True, op_level_per_channel=True)
qattention_name = "" if node.name == "" else node.name + "_quant"

4
onnxruntime/python/tools/quantization/operators/matmul.py
View file

@ -16,7 +16,7 @@ class MatMulInteger(QuantOperatorBase):
assert (node.op_type == "MatMul")
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer.quantize_inputs(node, [0, 1])
self.quantizer.quantize_inputs(node, [0, 1], reduce_range=True, op_level_per_channel=True)
matmul_integer_output = node.output[0] + "_output_quantized"
matmul_integer_name = node.name + "_quant" if node.name != "" else ""
@ -66,7 +66,7 @@ class QLinearMatMul(QuantOperatorBase):
assert (node.op_type == "MatMul")
(quantized_input_names, zero_point_names, scale_names, nodes) = \
self.quantizer.quantize_inputs(node, [0, 1])
self.quantizer.quantize_inputs(node, [0, 1], reduce_range=True, op_level_per_channel=True)
data_found, output_scale_name, output_zp_name, _, _ = \
self.quantizer._get_quantization_params(node.output[0])

22
onnxruntime/python/tools/quantization/quant_utils.py
View file

@ -109,11 +109,16 @@ def quantize_nparray(qType, arr, scale, zero_point, low=None, high=None):
return arr_fp32.astype(dtype)
def compute_scale_zp(rmin, rmax, qType, quantize_range):
def compute_scale_zp(rmin, rmax, qType, quantize_range, symmetric):
if qType == onnx_proto.TensorProto.INT8:
max_range = max(abs(rmin), abs(rmax))
scale = (float(max_range) * 2) / quantize_range if max_range > 0 else 1
zero_point = 0
if symmetric:
max_range = max(abs(rmin), abs(rmax))
scale = (float(max_range) * 2) / quantize_range if max_range > 0 else 1.0
zero_point = 0
else:
max_range = float(rmax) - float(rmin)
scale = float(max_range) / quantize_range if max_range > 0 else 1.0
zero_point = round((quantize_range / 2) - rmax / scale)
elif qType == onnx_proto.TensorProto.UINT8:
scale = (float(rmax) - rmin) / quantize_range if rmin != rmax else 1
zero_point = round((0 - rmin) / scale) # round to nearest integer
@ -123,11 +128,12 @@ def compute_scale_zp(rmin, rmax, qType, quantize_range):
return [zero_point, scale]
def quantize_data(data, quantize_range, qType):
def quantize_data(data, quantize_range, qType, symmetric=True):
'''
:parameter data: data to quantize
:parameter quantize_range: list of data to weight pack.
:parameter qType: data type to quantize to. Supported types UINT8 and INT8
:parameter symmetric: whether symmetric quantization is used or not. This is applied to INT8.
:return: minimum, maximum, zero point, scale, and quantized weights
To pack weights, we compute a linear transformation
- when data type == uint8 mode, from [rmin, rmax] -> [0, 2^{b-1}] and
@ -140,10 +146,10 @@ def quantize_data(data, quantize_range, qType):
S: scale
z: zero point
'''
rmin = min(min(data), 0)
rmax = max(max(data), 0)
rmin = min(min(data), 0.)
rmax = max(max(data), 0.)
zero_point, scale = compute_scale_zp(rmin, rmax, qType, quantize_range)
zero_point, scale = compute_scale_zp(rmin, rmax, qType, quantize_range, symmetric)
quantized_data = quantize_nparray(qType, numpy.asarray(data), scale, zero_point)
return rmin, rmax, zero_point, scale, quantized_data

5
onnxruntime/python/tools/transformers/fusion_attention.py
View file

@ -279,6 +279,11 @@ class FusionAttention(Fusion):
root_input = mul_before_layernorm.output[0]
else:
return
elif normalize_node.op_type == 'LayerNormalization':
children = input_name_to_nodes[root_input]
for child in children:
if child.op_type == "LayerNormalization":
root_input = child.output[0]
children = input_name_to_nodes[root_input]
children_types = [child.op_type for child in children]

63
onnxruntime/test/python/quantization/test_op_gemm.py
View file

@ -75,6 +75,52 @@ class TestOpGEMM(unittest.TestCase):
onnx.save(model, output_model_path)
def construct_model_attention_and_matmul(self, output_model_path):
# (input)
# |
# Attention
# |
# MatMul
# |
# (output)
input_name = 'input'
output_name = 'output'
initializers = []
def make_attention_node(input_name, weight_shape, weight_name, bias_shape, bias_name, output_name):
weight_data = np.random.normal(0, 0.1, weight_shape).astype(np.float32)
initializers.append(onnx.numpy_helper.from_array(weight_data, name=weight_name))
bias_data = np.random.normal(0, 0.1, bias_shape).astype(np.float32)
initializers.append(onnx.numpy_helper.from_array(bias_data, name=bias_name))
return onnx.helper.make_node('Attention', [input_name, weight_name, bias_name], [output_name])
def make_matmul_node(input_name, weight_shape, weight_name, output_name):
weight_data = np.random.normal(0, 0.1, weight_shape).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])
# make attention node
attention_output_name = "attention_output"
attention_node = make_attention_node(input_name, [10, 30], 'qkv.weight', [30], 'qkv.bias', attention_output_name)
attention_node.domain = "com.microsoft"
attention_node.attribute.extend([helper.make_attribute("num_heads", 5)])
# make matmul node
matmul_node = make_matmul_node(attention_output_name, [10, 10], 'matmul.weight', output_name)
# make graph
input_tensor = helper.make_tensor_value_info(input_name, TensorProto.FLOAT, [1, -1, 10])
output_tensor = helper.make_tensor_value_info(output_name, TensorProto.FLOAT, [1, -1, 10])
graph_name = 'attention_test'
graph = helper.make_graph([attention_node, matmul_node], graph_name,
[input_tensor], [output_tensor], initializer=initializers)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
model.ir_version = onnx.IR_VERSION
onnx.save(model, output_model_path)
def static_quant_test(self, model_fp32_path, model_int8_path):
data_reader = self.input_feeds(1, {'input': [5, 10]})
quantize_static(model_fp32_path,
@ -109,6 +155,12 @@ class TestOpGEMM(unittest.TestCase):
check_op_type_count(self, model_int8_path, **quant_nodes)
check_model_correctness(self, model_fp32_path, model_int8_path, {'input': np.random.rand(5,10).astype(np.float32)})
def dynamic_attention_quant_test(self, model_fp32_path, model_int8_path, per_channel, reduce_range):
quantize_dynamic(model_fp32_path, model_int8_path, per_channel=per_channel, reduce_range=reduce_range)
quant_nodes = {'QAttention' : 1, 'MatMulInteger' : 1}
check_op_type_count(self, model_int8_path, **quant_nodes)
check_model_correctness(self, model_fp32_path, model_int8_path, {'input': np.random.rand(1,5,10).astype(np.float32)})
def test_quantize_reshape(self):
np.random.seed(1)
model_fp32_path = 'gemm_fp32.onnx'
@ -119,5 +171,16 @@ class TestOpGEMM(unittest.TestCase):
self.static_quant_test_qdq(model_fp32_path, model_int8_path)
self.dynamic_quant_test(model_fp32_path, model_int8_path)
def test_quantize_attention(self):
np.random.seed(1)
model_fp32_path = 'attention_fp32.onnx'
model_int8_path = 'attention_fp32.quant.onnx'
self.construct_model_attention_and_matmul(model_fp32_path)
self.dynamic_attention_quant_test(model_fp32_path, model_int8_path, True, True)
self.dynamic_attention_quant_test(model_fp32_path, model_int8_path, True, False)
self.dynamic_attention_quant_test(model_fp32_path, model_int8_path, False, True)
self.dynamic_attention_quant_test(model_fp32_path, model_int8_path, False, False)
if __name__ == '__main__':
unittest.main()