[Torch] Support returning quantized weights and bias for BYOC use cas…

…es (#9135)

* [Torch] refactored the way is bias quantization done

* support returning 8bit weight

* add test

* add doc

* pylint

* return_int8_weight -> keep_quantized_weight

* fixed for dynamic linear case

* remove test function call

* simplifying

python/tvm/relay/frontend/pytorch.py

@@ -3713,6 +3713,7 @@ def from_pytorch(
     custom_convert_map=None,
     default_dtype="float32",
     use_parser_friendly_name=False,
+    keep_quantized_weight=False,
 ):
     """Load PyTorch model in the form of a scripted PyTorch model and convert into relay.
     The companion parameters will be handled automatically.
@@ -3745,6 +3746,16 @@ def from_pytorch(
         so a variable name like "dense.weight" cannot be parsed correctly.
         Use this option when you want to run the AnnotateSpans pass on the imported module.
 
+    keep_quantized_weight : bool
+        Return quantized weights and bias, rather than float ones. PyTorch stores quantized weights
+        in a custom format, so we cannot directly access 8 bit weights as Numpy arrays. We use
+        a PyTorch function to unpack quantized weights into float32 arrays and quantization
+        parameters. By default, we return float32 weights and rely on the QNN lowering and the
+        Relay constant folding pass to quantize weights at compile time. In BYOC use cases, however,
+        we cannot apply the constant folding pass on a QNN graph. If keep_quantized_weight is True,
+        we quantize weights in the frontend using a function that is equivalent to
+        qnn.op.quantize(...) operating on Numpy arrays.
+
     Returns
     -------
     mod : tvm.IRModule
@@ -3789,9 +3800,17 @@ def from_pytorch(
     # For quantized models
     quantized_ops = set(["aten::quantize_per_tensor", "quantized::linear_dynamic"])
     if len(quantized_ops.intersection(set(op_names))) > 0:
-        weight_quant_params = qnn_torch.get_weight_quant_params(script_module)
-        qnn_torch.add_input_quant_params_to_op_inputs(graph)
-        qnn_torch.add_quant_params_to_outputs(outputs, packed_param_map, weight_quant_params)
+        weight_quant_params = qnn_torch.get_weight_quant_params(
+            script_module, packed_param_map.values()
+        )
+        input_scales_for_bias = qnn_torch.add_input_quant_params_to_op_inputs(graph)
+        qnn_torch.add_quant_params_to_outputs(
+            outputs,
+            packed_param_map,
+            weight_quant_params,
+            input_scales_for_bias,
+            keep_quantized_weight,
+        )
         qnn_torch.add_quant_params(tvm_params, weight_quant_params)
         converter.update_convert_map(qnn_torch.convert_map)
 

python/tvm/relay/frontend/qnn_torch.py

@@ -32,16 +32,12 @@
 class QNNParam:
     """A placeholder for weight quantization parameters"""
 
-    def __init__(self, weight, bias, scale, zero_point, param_key):
-        param_prefix = param_key[: -len("._packed_params")]
-        self.weight_var = _expr.var(param_prefix + "_weight", shape=weight.shape)
+    def __init__(self, weight, bias, scale, zero_point):
         self.weight = weight
 
         if bias is not None:
-            self.bias_var = _expr.var(param_prefix + "_bias", shape=bias.shape)
             self.bias = bias.detach().numpy()
         else:
-            self.bias_var = None
             self.bias = None
 
         self.scale = _expr.const(scale)
@@ -55,10 +51,8 @@ class ConvPackedParam(QNNParam):
     together with weights and quantization parameters
     """
 
-    def __init__(
-        self, weight_np, bias, scale, zero_point, param_name, stride, padding, dilation, groups
-    ):
-        super().__init__(weight_np, bias, scale, zero_point, param_name)
+    def __init__(self, weight_np, bias, scale, zero_point, stride, padding, dilation, groups):
+        super().__init__(weight_np, bias, scale, zero_point)
         self.stride = stride
         self.padding = padding
         self.dilation = dilation
@@ -81,23 +75,21 @@ def _get_quant_params(qweight):
     return weight_np, scales, 0
 
 
-def make_qnn_param(param_name, qweight, bias):
+def make_qnn_param(qweight, bias):
     weight_np, scale, zero_point = _get_quant_params(qweight)
-    return QNNParam(weight_np, bias, scale, zero_point, param_name)
+    return QNNParam(weight_np, bias, scale, zero_point)
 
 
-def make_conv_packed_param(param_name, qweight, bias, packed_params):
+def make_conv_packed_param(qweight, bias, packed_params):
     weight_np, scale, zero_point = _get_quant_params(qweight)
     stride = packed_params.stride()
     padding = packed_params.padding()
     dilation = packed_params.dilation()
     groups = packed_params.groups()
-    return ConvPackedParam(
-        weight_np, bias, scale, zero_point, param_name, stride, padding, dilation, groups
-    )
+    return ConvPackedParam(weight_np, bias, scale, zero_point, stride, padding, dilation, groups)
 
 
-def get_weight_quant_params(script_module):
+def get_weight_quant_params(script_module, packed_param_names):
     """Retrive and unpack weight parameters from quantized modules"""
     import torch
 
@@ -114,6 +106,9 @@ def filter_func(named_module):
         key = name + "." + param_name
         state_dict = m.state_dict()
 
+        if key not in packed_param_names:
+            continue
+
         if len(state_dict) == 0 and not hasattr(m, param_name):
             # for v1.6 and above
             # This case seems to happen if a model is serialized
@@ -130,28 +125,87 @@ def filter_func(named_module):
 
         if "Conv" in m.original_name and len(state_dict) == 0:
             qweight, bias = torch.ops.quantized.conv2d_unpack(packed_params)
-            quant_params[key] = make_conv_packed_param(key, qweight, bias, packed_params)
+            quant_params[key] = make_conv_packed_param(qweight, bias, packed_params)
         elif "Conv" in m.original_name:
             qweight, bias = torch.ops.quantized.conv2d_unpack(packed_params)
-            quant_params[key] = make_qnn_param(key, qweight, bias)
+            quant_params[key] = make_qnn_param(qweight, bias)
         elif m.original_name == "LinearPackedParams":
             qweight, bias = torch.ops.quantized.linear_unpack(packed_params)
-            quant_params[key] = make_qnn_param(key, qweight, bias)
+            quant_params[key] = make_qnn_param(qweight, bias)
 
     return quant_params
 
 
-def add_quant_params_to_outputs(outputs, packed_param_map, quant_params):
+def quantize_numpy(weight, scale, zero_point, out_dtype_np):
+    iinfo = np.iinfo(out_dtype_np)
+    clip_min = iinfo.min
+    clip_max = iinfo.max
+    if len(scale.shape) > 0:
+        scale = np.reshape(scale, [weight.shape[0]] + [1] * (len(weight.shape) - 1))
+    transformed = zero_point + weight / scale
+    return np.clip(np.round(transformed), clip_min, clip_max).astype(out_dtype_np)
+
+
+def add_quant_params_to_outputs(
+    outputs, packed_param_map, quant_params, input_scales_for_bias, keep_quantized_weight=False
+):
     """
     Add quant params to outputs so that they can be referenced by other
     ops later. Weights are quantized here.
     """
     for node_name, packed_param_name in packed_param_map.items():
         qparam = quant_params[packed_param_name]
-        qweight = relay.qnn.op.quantize(
-            qparam.weight_var, qparam.scale, qparam.zero_point, out_dtype="int8", axis=0
-        )
-        params = [qweight, qparam.scale, qparam.zero_point, qparam.bias_var]
+        weight_scale = _get_numpy(qparam.scale)
+        param_prefix = packed_param_name[: -len("._packed_params")]
+
+        if keep_quantized_weight:
+            qparam.weight_var = _expr.var(
+                param_prefix + "_weight", shape=qparam.weight.shape, dtype="int8"
+            )
+            qparam.weight = quantize_numpy(
+                qparam.weight, weight_scale, _get_numpy(qparam.zero_point), np.int8
+            )
+            qweight = qparam.weight_var
+        else:
+            qparam.weight_var = _expr.var(
+                param_prefix + "_weight", shape=qparam.weight.shape, dtype="float32"
+            )
+            qweight = relay.qnn.op.quantize(
+                qparam.weight_var, qparam.scale, qparam.zero_point, out_dtype="int8", axis=0
+            )
+
+        if qparam.bias is not None:
+            float_bias_var = _expr.var(
+                param_prefix + "_bias", shape=qparam.bias.shape, dtype="float32"
+            )
+            if node_name not in input_scales_for_bias:
+                # This case is for dynamic quantization, where the input activation scale is
+                # unknown until runtime.
+                qparam.bias_var = float_bias_var
+                qbias = qparam.bias_var
+            elif keep_quantized_weight:
+                qparam.bias_var = _expr.var(
+                    param_prefix + "_bias", shape=qparam.bias.shape, dtype="int32"
+                )
+                qparam.bias = quantize_numpy(
+                    qparam.bias, input_scales_for_bias[node_name] * weight_scale, 0, np.int32
+                )
+                qbias = qparam.bias_var
+            else:
+                qparam.bias_var = float_bias_var
+                qbias = relay.qnn.op.quantize(
+                    qparam.bias_var,
+                    _expr.const(input_scales_for_bias[node_name] * weight_scale),
+                    _expr.const(0, "int32"),
+                    out_dtype="int32",
+                    axis=0,
+                )
+        else:
+            qbias = None
+
+        quant_params[packed_param_name] = qparam
+
+        params = [qweight, qparam.scale, qparam.zero_point, qbias]
 
         if isinstance(quant_params[packed_param_name], ConvPackedParam):
             params += [qparam.stride, qparam.padding, qparam.dilation, qparam.groups]
@@ -367,6 +421,8 @@ def add_input_quant_params_to_op_inputs(graph):
     need_input_quant_param = set(num_quantized_inputs.keys())
     need_input_quant_param.add("quantized::cat")
 
+    input_scales_for_bias = {}
+
     for node in graph.nodes():
         operator = node.kind()
         if operator not in need_input_quant_param:
@@ -401,6 +457,12 @@ def add_input_quant_params_to_op_inputs(graph):
             node.addInput(scale)
             node.addInput(zp)
 
+        if "conv2d" in operator or "linear" in operator:
+            # This is required for quantizing the bias
+            input_scales_for_bias[node.inputsAt(1).debugName()] = scale.node().f("value")
+
+    return input_scales_for_bias
+
 
 def add_quant_params(params, quant_params):
     """Add quant parameters to TVM param map"""
@@ -478,10 +540,7 @@ def _do_bias_and_requantize(
     # Instead, the torch way requires rounding of activation at runtime
 
     if bias is not None:
-        qbias = relay.qnn.op.quantize(
-            bias, requant_input_scale, _expr.const(0, "int32"), out_dtype="int32", axis=0
-        )
-        requantize_input = _op.nn.bias_add(output, qbias)
+        requantize_input = _op.nn.bias_add(output, bias)
     else:
         requantize_input = output
 

tests/python/frontend/pytorch/qnn_test.py

@@ -40,9 +40,15 @@ def torch_version_check():
     return version.parse(torch.__version__) > version.parse("1.4.0")
 
 
-def get_tvm_runtime(script_module, input_name, ishape):
+def get_tvm_runtime(script_module, input_name, ishape, keep_quantized_weight=False):
     input_shapes = [(input_name, ishape)]
-    mod, params = relay.frontend.from_pytorch(script_module, input_shapes)
+    mod, params = relay.frontend.from_pytorch(
+        script_module, input_shapes, keep_quantized_weight=keep_quantized_weight
+    )
+
+    if keep_quantized_weight:
+        for p in params.values():
+            assert p.dtype in ["int8", "int32"]
 
     with tvm.transform.PassContext(opt_level=3):
         # test on only cpu for now, torch cannot run quant models on cuda
@@ -609,3 +615,36 @@ def test_qnn_mergecomposite():
 
     input_name = "image"
     run_qnn_mergecomposite(script_module, input_name, inp.shape)
+
+
+def test_keep_quantized_weight():
+    qmodules = []
+
+    for per_channel in [False, True]:
+        qmodules += [
+            ((1, 3, 224, 224), ConvBn(), per_channel),
+            ((16, 16), Linear(), per_channel),
+        ]
+
+    for (ishape, raw_module, per_channel) in qmodules:
+        raw_module.eval()
+        inp = torch.rand(ishape)
+
+        quantize_model(raw_module, inp, per_channel=per_channel)
+        script_module = torch.jit.trace(raw_module, inp).eval()
+
+        input_name = "input"
+
+        runtime = get_tvm_runtime(script_module, input_name, ishape, keep_quantized_weight=False)
+        runtime.set_input(input_name, inp.numpy().copy())
+        runtime.run()
+        tvm_result = runtime.get_output(0).numpy()
+
+        runtime_int8_weight = get_tvm_runtime(
+            script_module, input_name, ishape, keep_quantized_weight=True
+        )
+        runtime_int8_weight.set_input(input_name, inp.numpy().copy())
+        runtime_int8_weight.run()
+        tvm_result_int8_weight = runtime_int8_weight.get_output(0).numpy()
+
+        tvm.testing.assert_allclose(tvm_result, tvm_result_int8_weight)