From a009e17d2b3b8eea134fb6d05840a36a7d0770e8 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 17 May 2025 19:46:47 +0300 Subject: [PATCH] NNCF use per token input quantization with int8 matmul --- modules/model_quant_nncf.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/modules/model_quant_nncf.py b/modules/model_quant_nncf.py index 9e8818330..3d8fb4624 100644 --- a/modules/model_quant_nncf.py +++ b/modules/model_quant_nncf.py @@ -461,9 +461,11 @@ def unpack_int4(packed_tensor: torch.Tensor, shape: torch.Size, dtype: Optional[ def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.ByteTensor, torch.FloatTensor]: - input_scale = torch.div(input.abs().max(), 127) - input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(torch.int8).flatten(0,-2) - scale = torch.mul(input_scale, scale) + input_scale = torch.div(input.abs().max(dim=-1).values, 127).unsqueeze(-1) + input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(torch.int8).flatten(0,-2).contiguous() + scale = torch.mul(input_scale, scale).flatten(0,-2).contiguous() + if scale.dtype == torch.float16: # fp16 will overflow + scale = scale.to(dtype=torch.float32) return input, scale @@ -485,7 +487,7 @@ def int8_matmul( class linear_forward_int8_matmul(): def __func__(self, input) -> torch.FloatTensor: if self.pre_ops["0"].skip_int8_matmul: - return torch.nn.Linear.forward(self, input) + return torch.nn.functional.linear(input, self.weight, self.bias) result = int8_matmul(input, self.weight, self.pre_ops["0"].scale, getattr(self.pre_ops["0"], "compressed_weight_shape", None)) if self.bias is not None: result.add_(self.bias)