File size: 6,054 Bytes
610d9c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | import logging
import numpy as np # noqa: F401
import onnx
from ..quant_utils import (
TENSOR_NAME_QUANT_SUFFIX,
QuantizedValue,
QuantizedValueType,
attribute_to_kwarg,
find_by_name, # noqa: F401
get_mul_node, # noqa: F401
ms_domain,
)
from .base_operator import QuantOperatorBase # noqa: F401
from .matmul import QOpMatMul
from .qdq_base_operator import QDQOperatorBase
def is_B_transposed(gemm_node): # noqa: N802
transB_attribute = [attr for attr in gemm_node.attribute if attr.name == "transB"] # noqa: N806
if transB_attribute:
return onnx.helper.get_attribute_value(transB_attribute[0]) > 0
return False
def get_beta(gemm_node):
beta_attribute = [attr for attr in gemm_node.attribute if attr.name == "beta"]
if beta_attribute:
return onnx.helper.get_attribute_value(beta_attribute[0])
return 1.0
def set_default_beta(gemm_node):
beta_attribute = [attr for attr in gemm_node.attribute if attr.name == "beta"]
if beta_attribute:
beta_attribute[0].f = 1.0
return 1.0
class QLinearGemm(QOpMatMul):
def __init__(self, onnx_quantizer, onnx_node):
super().__init__(onnx_quantizer, onnx_node)
def quantize(self):
node = self.node
assert node.op_type == "Gemm"
(
data_found,
output_scale_name,
output_zp_name,
_,
_,
) = self.quantizer._get_quantization_params(node.output[0])
if self.quantizer.is_input_a_initializer(node.input[1]) and self.quantizer.is_per_channel():
(
quantized_input_names,
zero_point_names,
scale_names,
nodes,
) = self.quantizer.quantize_activation(node, [0])
quant_weight_tuple = self.quantizer.quantize_weight_per_channel(
node.input[1],
self.quantizer.weight_qType,
0 if is_B_transposed(node) else 1,
)
quantized_input_names.append(quant_weight_tuple[0])
zero_point_names.append(quant_weight_tuple[1])
scale_names.append(quant_weight_tuple[2])
else:
# Get Quantized from both activation(input[0]) and weight(input[1])
(
quantized_input_names,
zero_point_names,
scale_names,
nodes,
) = self.quantizer.quantize_activation(node, [0])
(
quantized_input_names_weight,
zero_point_names_weight,
scale_names_weight,
nodes_weight,
) = self.quantizer.quantize_weight(node, [1], reduce_range=self.quantizer.reduce_range)
quantized_input_names.extend(quantized_input_names_weight)
zero_point_names.extend(zero_point_names_weight)
scale_names.extend(scale_names_weight)
nodes.extend(nodes_weight)
if not data_found or quantized_input_names is None:
return super().quantize()
quantized_bias_name = ""
if len(node.input) == 3:
if not self.quantizer.is_input_a_initializer(node.input[2]):
return super().quantize()
# Note: if the quantized type is float 8, the bias is converted into float 16.
# cublasLtMatMul only supports (b)float16 or float32 bias.
quantized_bias_name = self.quantizer.quantize_bias_static(
node.input[2], node.input[0], node.input[1], get_beta(self.node)
)
qgemm_output = node.output[0] + TENSOR_NAME_QUANT_SUFFIX
qgemm_name = node.name + "_quant" if node.name else ""
kwargs = {}
for attribute in node.attribute:
if attribute.name != "beta":
kwargs.update(attribute_to_kwarg(attribute))
kwargs["domain"] = ms_domain
# generate input
qgemm_inputs = []
for i in range(2):
qgemm_inputs.extend([quantized_input_names[i], scale_names[i], zero_point_names[i]])
qgemm_inputs.extend([quantized_bias_name, output_scale_name, output_zp_name])
qgemm_node = onnx.helper.make_node("QGemm", qgemm_inputs, [qgemm_output], qgemm_name, **kwargs)
nodes.append(qgemm_node)
# Create an entry for this quantized value
q_output = QuantizedValue(
node.output[0],
qgemm_output,
output_scale_name,
output_zp_name,
QuantizedValueType.Input,
node_type=node.op_type,
node_qtype=self.quantizer.weight_qType,
)
self.quantizer.quantized_value_map[node.output[0]] = q_output
self.quantizer.new_nodes += nodes
class QDQGemm(QDQOperatorBase):
def __init__(self, onnx_quantizer, onnx_node):
super().__init__(onnx_quantizer, onnx_node)
def quantize(self):
node = self.node
assert node.op_type == "Gemm"
self.quantizer.quantize_activation_tensor(node.input[0])
if not self.disable_qdq_for_node_output:
self.quantizer.quantize_activation_tensor(node.output[0])
is_weight_per_channel, weight_axis = self.quantizer.is_tensor_per_channel(
node.input[1], default_axis=0 if is_B_transposed(node) else 1
)
if is_weight_per_channel:
self.quantizer.quantize_weight_tensor_per_channel(node.input[1], weight_axis)
else:
self.quantizer.quantize_weight_tensor(node.input[1])
if len(node.input) == 3:
if self.quantizer.is_input_a_initializer(node.input[2]):
self.quantizer.quantize_bias_tensor(
node.name, node.input[2], node.input[0], node.input[1], get_beta(self.node)
)
set_default_beta(self.node)
else:
logging.warning(
f"Bias of Gemm node '{self.node.name}' is not constant. Please exclude this node for better performance."
)
|