diff --git a/modules/model_quant_sdnq.py b/modules/model_quant_sdnq.py index 1281b05e4..dbf0cfa4f 100644 --- a/modules/model_quant_sdnq.py +++ b/modules/model_quant_sdnq.py @@ -57,21 +57,30 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz if layer_class_name in conv_types: if not quant_conv: return layer + if dtype_dict[weights_dtype]["num_bits"] < 4: + weights_dtype = "uint4" is_conv_type = True reduction_axes = 1 output_channel_size, channel_size = layer.weight.shape[:2] + group_channel_size = channel_size // layer.groups use_quantized_matmul = False - if dtype_dict[weights_dtype]["num_bits"] < 4: - weights_dtype = "uint4" + if shared.opts.sdnq_use_quantized_matmul_conv: + use_quantized_matmul = dtype_dict[weights_dtype]["is_integer"] and weights_dtype in quantized_matmul_dtypes and group_channel_size >= 32 and output_channel_size >= 32 + #if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]: + # use_quantized_matmul = output_channel_size % 16 == 0 and group_channel_size % 16 == 0 + # use_tensorwise_fp8_matmul = torch_version < 2.5 or devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and torch_version <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9)) + if use_quantized_matmul: + result_shape = layer.weight.shape + layer.weight.data = layer.weight.reshape(output_channel_size, -1) elif layer_class_name in conv_transpose_types: if not quant_conv: return layer + if dtype_dict[weights_dtype]["num_bits"] < 4: + weights_dtype = "uint4" is_conv_transpose_type = True reduction_axes = 0 channel_size, output_channel_size = layer.weight.shape[:2] use_quantized_matmul = False - if dtype_dict[weights_dtype]["num_bits"] < 4: - weights_dtype = "uint4" else: is_linear_type = True reduction_axes = -1 @@ -193,7 +202,10 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz else: layer.forward = quantized_linear_forward elif is_conv_type: - layer.forward = quantized_conv_forward + if use_quantized_matmul: + layer.forward = quantized_conv2d_forward_int8_matmul + else: + layer.forward = quantized_conv_forward elif is_conv_transpose_type: if layer_class_name.endswith("1d"): layer.forward = quantized_conv_transpose_1d_forward @@ -393,7 +405,7 @@ def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch. return input, scale -def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.ByteTensor, torch.FloatTensor]: +def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor, flatten: bool = True) -> Tuple[torch.ByteTensor, torch.FloatTensor]: input = input.flatten(0,-2).contiguous() input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 127) input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(torch.int8) @@ -454,6 +466,54 @@ def int8_matmul( return result +def conv2d_int8_matmul( + input: torch.FloatTensor, + weight: torch.ByteTensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, + result_shape: torch.Size, + compressed_weight_shape: torch.Size, + weights_dtype: str, + reversed_padding_repeated_twice: List[int], + padding_mode: str, groups: int, + stride_h: int, stride_w: int, + padding_h: int, padding_w: int, + dilation_h: int, dilation_w: int, +) -> torch.FloatTensor: + return_dtype = input.dtype + batch_size, _, H_in, W_in = input.shape + C_out, _, K_h, K_w = result_shape + W_out = (W_in + 2 * padding_w - dilation_w * (K_w - 1) - 1) // stride_w + 1 + H_out = (H_in + 2 * padding_h - dilation_h * (K_h - 1) - 1) // stride_h + 1 + mm_output_shape = (batch_size, H_out, W_out, C_out) + + if compressed_weight_shape is not None: + weight = unpack_int_symetric(weight, compressed_weight_shape, weights_dtype, dtype=torch.int8, transpose=True) + if padding_mode != "zeros": + input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode) + padding_h = padding_w = 0 + + input, scale = quantize_int8_matmul_input( + torch.nn.functional.unfold( + input, kernel_size=(K_h, K_w), padding=(padding_h, padding_w), stride=(stride_h, stride_w), dilation=(dilation_h, dilation_w) + ).transpose(1,2), + scale, + ) + + if groups == 1: + result = decompress_symmetric(torch._int_mm(input, weight), scale, return_dtype, mm_output_shape) + else: + weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) + input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) + result = [] + for i in range(groups): + result.append(torch._int_mm(input[i], weight[i])) + result = decompress_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape) + if bias is not None: + result.add_(bias) + return result.permute(0,3,1,2) + + def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: if torch.numel(input) / input.shape[-1] < 32: return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias) @@ -472,6 +532,36 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc return int8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale, getattr(self.sdnq_decompressor, "compressed_weight_shape", None), self.sdnq_decompressor.weights_dtype) +def quantized_conv2d_forward_int8_matmul(self, input) -> torch.FloatTensor: + if isinstance(self.stride, int): + stride_h = stride_w = self.stride + else: + stride_h, stride_w = self.stride + + if isinstance(self.padding, int): + padding_h = padding_w = self.padding + else: + padding_h, padding_w = self.padding + + if isinstance(self.dilation, int): + dilation_h = dilation_w = self.dilation + else: + dilation_h, dilation_w = self.dilation + + return conv2d_int8_matmul( + input, self.weight, self.bias, + self.sdnq_decompressor.scale, + self.sdnq_decompressor.result_shape, + getattr(self.sdnq_decompressor, "compressed_weight_shape", None), + self.sdnq_decompressor.weights_dtype, + self._reversed_padding_repeated_twice, + self.padding_mode, self.groups, + stride_h, stride_w, + padding_h, padding_w, + dilation_h, dilation_w, + ) + + def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor: return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight), self.bias) @@ -840,6 +930,7 @@ if shared.opts.sdnq_decompress_compile: fp8_matmul = torch.compile(fp8_matmul, fullgraph=True) fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True) int8_matmul = torch.compile(int8_matmul, fullgraph=True) + conv2d_int8_matmul = torch.compile(conv2d_int8_matmul, fullgraph=True) except Exception as e: shared.log.warning(f"Quantization: type=sdnq Decompress using torch.compile is not available: {e}") decompress_asymmetric_compiled = decompress_asymmetric diff --git a/modules/shared.py b/modules/shared.py index a1c089c48..6fdd80ad3 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -521,10 +521,11 @@ options_templates.update(options_section(("quantization", "Quantization Settings "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ["int8", "int6", "uint4", "float8_e4m3fn", "uint8", "uint6", "int4", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "int2", "uint2", "uint1"], "visible": native}), "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": native}), "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize the convolutional layers", gr.Checkbox, {"visible": native}), - "sdnq_decompress_fp32": OptionInfo(False, "Decompress using full precision", gr.Checkbox, {"visible": native}), "sdnq_decompress_compile": OptionInfo(devices.has_triton(), "Decompress using torch.compile", gr.Checkbox, {"visible": native}), "sdnq_use_quantized_matmul": OptionInfo(False, "Use Quantized MatMul", gr.Checkbox, {"visible": native}), + "sdnq_use_quantized_matmul_conv": OptionInfo(False, "Use Quantized MatMul with convolutional layers", gr.Checkbox, {"visible": native}), "sdnq_quantize_with_gpu": OptionInfo(True, "Quantize with the GPU", gr.Checkbox, {"visible": native}), + "sdnq_decompress_fp32": OptionInfo(False, "Decompress using full precision", gr.Checkbox, {"visible": native}), "sdnq_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}), "bnb_quantization_sep": OptionInfo("

BitsAndBytes

", "", gr.HTML),