Updated QOp model to fuse SmoothQuant scales with input quantization

math_model.py

@@ -47,9 +47,9 @@ class QuantLinear(nn.Module):
         #  - multiply this sum with every weight zero-point (e.g., `torch.sum(quant_input, dim=-1) * self.weight_zp`
         #  - Subtract from previous output (e.g., `quant_output -= torch.sum(quant_input, dim=-1) * self.weight_zp`)
         #  - All other code is just to make sure the broadcasting semantics work correctly
-        scaled_x = x * self.mul_factor
         quant_weight = quantize(self.linear.weight, self.weight_scale, self.weight_zp, is_asym=True).to(torch.uint8)
-        quant_input = quantize(scaled_x, self.input_scale, self.input_zp, is_asym=False).to(torch.int8)
+        fused_input_scale = self.input_scale / self.mul_factor # Fuse SmoothQuant and input scales, can be computed offline
+        quant_input = quantize(x, fused_input_scale, self.input_zp, is_asym=False).to(torch.int8)
         quant_output = torch.nn.functional.linear(quant_input.to(torch.float32), quant_weight.to(torch.float32), None).to(torch.int32) # Convert inputs to FP32 to avoid F.linear quantizing the output to int8
         correction = torch.sum(quant_input, dim=-1, keepdim=True).to(torch.int32) * (-self.weight_zp).to(torch.uint8).view([1]*(quant_input.ndim-1) + [self.weight_zp.nelement()]) # Correct for weight zero-point
         quant_output = quant_output + correction
@@ -103,13 +103,13 @@ class QuantConv2d(nn.Module):
         #  - multiply this sum with every weight zero-point (e.g., `sum * self.weight_zp`
         #  - Subtract from previous output (e.g., `quant_output -= sum * self.weight_zp`)
         #  - All other code is just to make sure the broadcasting semantics work correctly
-        scaled_x = x * self.mul_factor
         quant_weight = quantize(self.conv2d.weight, self.weight_scale, self.weight_zp, is_asym=True).to(torch.uint8)
         b_shape = list(quant_weight.shape) # Used for weight zero-point correction
         b_shape[0] = 1 # Used for weight zero-point correction
         weight_cat = torch.ones((1,1,1,1)).broadcast_to(b_shape).to(torch.uint8) # Used for weight zero-point correction
         quant_weight = torch.cat((quant_weight,weight_cat),dim=0).to(torch.uint8) # Create extra output channel, used for weight zero-point correction
-        quant_input = quantize(scaled_x, self.input_scale, self.input_zp, is_asym=False).to(torch.int8)
+        fused_input_scale = self.input_scale / self.mul_factor # Fuse SmoothQuant and input scales, can be computed offline
+        quant_input = quantize(x, fused_input_scale, self.input_zp, is_asym=False).to(torch.int8)
         quant_output = torch.nn.functional.conv2d(quant_input.to(torch.float32), quant_weight.to(torch.float32), None).to(torch.int32) # Convert inputs to FP32 to avoid F.conv2d quantizing the output to int8
         correction = quant_output[:,-1,:,:] * (-self.weight_zp).to(torch.uint8).view([1, self.weight_zp.nelement()] +  [1]*(quant_output.ndim-2)) # Correct zero-point for weight
         quant_output = quant_output[:,:-1,:,:] + correction