Remove potential overflow / saturation error.

math_model.py

@@ -52,8 +52,8 @@ class QuantLinear(nn.Module):
         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) * (-weight_zp_int8).to(torch.int8).view([1]*(quant_input.ndim-1) + [self.weight_zp.nelement()]) # Correct for weight zero-point, possible overflow / saturation on (-weight_zp_int8).to(torch.int8)
-        quant_output = quant_output + correction
+        correction = torch.sum(quant_input, dim=-1, keepdim=True).to(torch.int32) * weight_zp_int8.to(torch.int8).view([1]*(quant_input.ndim-1) + [self.weight_zp.nelement()]) # Correct for weight zero-point
+        quant_output = quant_output - correction
         output = dequantize(quant_output, (self.weight_scale * self.input_scale).view([1]*(quant_output.ndim-1) + [(self.weight_scale * self.input_scale).nelement()]), 0.0)
         output += self.linear.bias
         return output
@@ -113,8 +113,8 @@ class QuantConv2d(nn.Module):
         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,:,:] * (-weight_zp_int8).to(torch.int8).view([1, self.weight_zp.nelement()] +  [1]*(quant_output.ndim-2)) # Correct zero-point for weight, possible overflow / saturation on (-weight_zp_int8).to(torch.int8)
-        quant_output = quant_output[:,:-1,:,:] + correction
+        correction = quant_output[:,-1,:,:] * weight_zp_int8.to(torch.int8).view([1, self.weight_zp.nelement()] +  [1]*(quant_output.ndim-2)) # Correct zero-point for weight
+        quant_output = quant_output[:,:-1,:,:] - correction
         output = dequantize(quant_output, (self.weight_scale * self.input_scale).view([1, (self.weight_scale * self.input_scale).nelement()] + [1]*(quant_output.ndim-2)), 0.0)
         output += self.conv2d.bias.view([1, self.conv2d.bias.nelement()] + [1]*(quant_output.ndim-2))
         return output