diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index 0343d7b14..47b408446 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -47,6 +47,11 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n self.network_weights_backup = weight.clone().to(devices.cpu) if hasattr(self, "sdnq_dequantizer"): self.sdnq_dequantizer_backup = self.sdnq_dequantizer.to(devices.cpu) + self.sdnq_scale_backup = self.scale.clone().to(devices.cpu) + if self.zero_point is not None: + self.sdnq_zero_point_backup = self.zero_point.clone().to(devices.cpu) + else: + self.sdnq_zero_point_backup = None if bias_backup is None: if getattr(self, 'bias', None) is not None: @@ -80,9 +85,9 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn. try: t0 = time.time() if hasattr(self, "sdnq_dequantizer_backup"): - weight = self.sdnq_dequantizer_backup.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_dequantizer_backup.use_quantized_matmul) + weight = self.sdnq_dequantizer_backup.to(devices.device)(self.weight.to(devices.device), self.sdnq_scale_backup.to(devices.device), self.sdnq_zero_point_backup.to(devices.device) if self.sdnq_zero_point_backup is not None else None, skip_quantized_matmul=self.sdnq_dequantizer_backup.use_quantized_matmul) elif hasattr(self, "sdnq_dequantizer"): - weight = self.sdnq_dequantizer.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul) + weight = self.sdnq_dequantizer.to(devices.device)(self.weight.to(devices.device), self.scale.to(devices.device), self.zero_point.to(devices.device) if self.zero_point is not None else None, skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul) else: weight = self.weight.to(devices.device) # must perform calc on gpu due to performance updown, ex_bias = module.calc_updown(weight) @@ -142,16 +147,18 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G try: from modules.sdnq import sdnq_quantize_layer if hasattr(self, "sdnq_dequantizer_backup"): - sdnq_dequantizer = self.sdnq_dequantizer_backup.to(devices.device) - else: - sdnq_dequantizer = self.sdnq_dequantizer.to(devices.device) - dequant_weight = sdnq_dequantizer(model_weights.to(devices.device), skip_quantized_matmul=sdnq_dequantizer.use_quantized_matmul) + weights_dtype = self.sdnq_dequantizer_backup.weights_dtype + dequant_weight = self.sdnq_dequantizer_backup.to(devices.device)(model_weights.to(devices.device), self.sdnq_scale_backup.to(devices.device), self.sdnq_zero_point_backup.to(devices.device) if self.sdnq_zero_point_backup is not None else None, skip_quantized_matmul=self.sdnq_dequantizer_backup.use_quantized_matmul) + elif hasattr(self, "sdnq_dequantizer"): + weights_dtype = self.sdnq_dequantizer.weights_dtype + dequant_weight = self.sdnq_dequantizer.to(devices.device)(model_weights.to(devices.device), self.scale.to(devices.device), self.zero_point.to(devices.device) if self.zero_point is not None else None, skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul) + new_weight = dequant_weight.to(devices.device, dtype=torch.float32) + lora_weights.to(devices.device, dtype=torch.float32) self.weight = torch.nn.Parameter(new_weight, requires_grad=False) - self.sdnq_dequantizer = None + del self.sdnq_dequantizer, self.scale, self.zero_point self = sdnq_quantize_layer( self, - sdnq_dequantizer.weights_dtype, + weights_dtype=weights_dtype, torch_dtype=devices.dtype, group_size=shared.opts.sdnq_quantize_weights_group_size, quant_conv=shared.opts.sdnq_quantize_conv_layers, @@ -231,7 +238,12 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn self.weight = torch.nn.Parameter(weights_backup.to(device), requires_grad=False) if hasattr(self, "sdnq_dequantizer_backup"): self.sdnq_dequantizer = self.sdnq_dequantizer_backup.to(device) - del self.sdnq_dequantizer_backup + self.scale = torch.nn.Parameter(self.sdnq_scale_backup.to(device), requires_grad=False) + if self.sdnq_zero_point_backup is not None: + self.zero_point = torch.nn.Parameter(self.sdnq_zero_point_backup.to(device), requires_grad=False) + else: + self.zero_point = None + del self.sdnq_dequantizer_backup, self.sdnq_scale_backup, self.sdnq_zero_point_backup if bias_backup is not None: self.bias = None diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index a84f48f2e..16fe6ec47 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -179,9 +179,15 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz if not use_tensorwise_fp8_matmul and not dtype_dict[weights_dtype]["is_integer"]: scale = scale.to(torch.float32) + scale = scale.to(return_device, non_blocking=non_blocking) + layer.scale = torch.nn.Parameter(scale, requires_grad=False) + if zero_point is not None: + zero_point = zero_point.to(return_device, non_blocking=non_blocking) + layer.zero_point = torch.nn.Parameter(zero_point, requires_grad=False) + else: + layer.zero_point = None + layer.sdnq_dequantizer = dequantizer_dict[weights_dtype]( - scale=scale, - zero_point=zero_point, quantized_weight_shape=layer.weight.shape, result_dtype=torch_dtype, result_shape=result_shape, @@ -191,7 +197,6 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz re_quantize_for_matmul=re_quantize_for_matmul, ) layer.weight.data = layer.sdnq_dequantizer.pack_weight(layer.weight).to(return_device, non_blocking=non_blocking) - layer.sdnq_dequantizer = layer.sdnq_dequantizer.to(return_device, non_blocking=non_blocking) layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, dtype_dict[weights_dtype]["is_integer"], use_tensorwise_fp8_matmul) layer.forward = layer.forward.__get__(layer, layer.__class__) diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index be68ab96d..d82849177 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -73,8 +73,6 @@ def re_quantize_matmul_packed_int_symmetric(weight: torch.ByteTensor, scale: tor class AsymmetricWeightsDequantizer(torch.nn.Module): def __init__( self, - scale: torch.FloatTensor, - zero_point: torch.FloatTensor, result_dtype: torch.dtype, result_shape: torch.Size, original_shape: torch.Size, @@ -89,23 +87,20 @@ class AsymmetricWeightsDequantizer(torch.nn.Module): self.re_quantize_for_matmul = True self.result_dtype = result_dtype self.result_shape = result_shape - self.register_buffer("scale", scale) - self.register_buffer("zero_point", zero_point) def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]) - def re_quantize_matmul(self, weight, **kwargs): # pylint: disable=unused-argument - return re_quantize_matmul_asymmetric_compiled(weight, self.scale, self.zero_point, self.result_shape) + def re_quantize_matmul(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument + return re_quantize_matmul_asymmetric_compiled(weight, scale, zero_point, self.result_shape) - def forward(self, weight, **kwargs): # pylint: disable=unused-argument - return dequantize_asymmetric_compiled(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape) + def forward(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument + return dequantize_asymmetric_compiled(weight, scale, zero_point, self.result_dtype, self.result_shape) class SymmetricWeightsDequantizer(torch.nn.Module): def __init__( self, - scale: torch.FloatTensor, result_dtype: torch.dtype, result_shape: torch.Size, original_shape: torch.Size, @@ -121,24 +116,21 @@ class SymmetricWeightsDequantizer(torch.nn.Module): self.re_quantize_for_matmul = re_quantize_for_matmul self.result_dtype = result_dtype self.result_shape = result_shape - self.register_buffer("scale", scale) def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]) - def re_quantize_matmul(self, weight, **kwargs): # pylint: disable=unused-argument - return re_quantize_matmul_symmetric_compiled(weight, self.scale, self.result_shape) + def re_quantize_matmul(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument + return re_quantize_matmul_symmetric_compiled(weight, scale, self.result_shape) - def forward(self, weight, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument + def forward(self, weight, scale, zero_point, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument skip_quantized_matmul = skip_quantized_matmul and not self.re_quantize_for_matmul - return dequantize_symmetric_compiled(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul) + return dequantize_symmetric_compiled(weight, scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul) class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module): def __init__( self, - scale: torch.FloatTensor, - zero_point: torch.FloatTensor, quantized_weight_shape: torch.Size, result_dtype: torch.dtype, result_shape: torch.Size, @@ -155,23 +147,20 @@ class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module): self.quantized_weight_shape = quantized_weight_shape self.result_dtype = result_dtype self.result_shape = result_shape - self.register_buffer("scale", scale) - self.register_buffer("zero_point", zero_point) def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: return pack_int_asymetric(weight, self.weights_dtype) - def re_quantize_matmul(self, weight, **kwargs): # pylint: disable=unused-argument - return re_quantize_matmul_packed_int_asymmetric_compiled(weight, self.scale, self.zero_point, self.quantized_weight_shape, self.result_shape, self.weights_dtype) + def re_quantize_matmul(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument + return re_quantize_matmul_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.result_shape, self.weights_dtype) - def forward(self, weight, **kwargs): # pylint: disable=unused-argument - return dequantize_packed_int_asymmetric_compiled(weight, self.scale, self.zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype) + def forward(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument + return dequantize_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype) class PackedINTSymmetricWeightsDequantizer(torch.nn.Module): def __init__( self, - scale: torch.FloatTensor, quantized_weight_shape: torch.Size, result_dtype: torch.dtype, result_shape: torch.Size, @@ -189,17 +178,16 @@ class PackedINTSymmetricWeightsDequantizer(torch.nn.Module): self.quantized_weight_shape = quantized_weight_shape self.result_dtype = result_dtype self.result_shape = result_shape - self.register_buffer("scale", scale) def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: return pack_int_symetric(weight, self.weights_dtype) - def re_quantize_matmul(self, weight, **kwargs): # pylint: disable=unused-argument - return re_quantize_matmul_packed_int_symmetric_compiled(weight, self.scale, self.quantized_weight_shape, self.result_shape, self.weights_dtype) + def re_quantize_matmul(self, weight, scale, zero_point, **kwargs): # pylint: disable=unused-argument + return re_quantize_matmul_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.result_shape, self.weights_dtype) - def forward(self, weight, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument + def forward(self, weight, scale, zero_point, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument skip_quantized_matmul = skip_quantized_matmul and not self.re_quantize_for_matmul - return dequantize_packed_int_symmetric_compiled(weight, self.scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul) + return dequantize_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul) dequantizer_dict = { diff --git a/modules/sdnq/layers/conv/conv_fp8.py b/modules/sdnq/layers/conv/conv_fp8.py index 7a83bb3a8..47f214937 100644 --- a/modules/sdnq/layers/conv/conv_fp8.py +++ b/modules/sdnq/layers/conv/conv_fp8.py @@ -58,11 +58,10 @@ def conv_fp8_matmul( def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor: if torch.numel(input) / input.shape[2] < 32: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) return conv_fp8_matmul( - input, self.weight, self.bias, - self.sdnq_dequantizer.scale, + input, self.weight, self.bias, self.scale, self.sdnq_dequantizer.result_shape, self._reversed_padding_repeated_twice, self.padding_mode, conv_type, diff --git a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py index e3b005622..592a512ee 100644 --- a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py +++ b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py @@ -53,11 +53,10 @@ def conv_fp8_matmul_tensorwise( def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTensor: if torch.numel(input) / input.shape[2] < 32: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) return conv_fp8_matmul_tensorwise( - input, self.weight, self.bias, - self.sdnq_dequantizer.scale, + input, self.weight, self.bias, self.scale, self.sdnq_dequantizer.result_shape, self._reversed_padding_repeated_twice, self.padding_mode, conv_type, diff --git a/modules/sdnq/layers/conv/conv_int8.py b/modules/sdnq/layers/conv/conv_int8.py index 260acd170..0ae26b8f5 100644 --- a/modules/sdnq/layers/conv/conv_int8.py +++ b/modules/sdnq/layers/conv/conv_int8.py @@ -57,14 +57,14 @@ def conv_int8_matmul( def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor: if torch.numel(input) / input.shape[2] < 32: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) if self.sdnq_dequantizer.re_quantize_for_matmul: - weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight) + weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point) quantized_weight_shape = None else: weight = self.weight - scale = self.sdnq_dequantizer.scale + scale = self.scale quantized_weight_shape = getattr(self.sdnq_dequantizer, "quantized_weight_shape", None) return conv_int8_matmul( input, weight, self.bias, diff --git a/modules/sdnq/layers/conv/forward.py b/modules/sdnq/layers/conv/forward.py index 44c90b061..69237ea68 100644 --- a/modules/sdnq/layers/conv/forward.py +++ b/modules/sdnq/layers/conv/forward.py @@ -75,19 +75,19 @@ def process_conv_input(conv_type, input, reversed_padding_repeated_twice, paddin def quantized_conv_forward(self, input) -> torch.FloatTensor: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight), self.bias) + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point), self.bias) def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation) - return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) + return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation) - return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) + return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation) - return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) + return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) diff --git a/modules/sdnq/layers/linear/forward.py b/modules/sdnq/layers/linear/forward.py index 6977d6566..fc5b2ce60 100644 --- a/modules/sdnq/layers/linear/forward.py +++ b/modules/sdnq/layers/linear/forward.py @@ -15,4 +15,4 @@ def check_mats(input: torch.Tensor, weight: torch.Tensor) -> Tuple[torch.Tensor, def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor: - return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias) + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point), self.bias) diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py index d1845b1e7..3f2fb407c 100644 --- a/modules/sdnq/layers/linear/linear_fp8.py +++ b/modules/sdnq/layers/linear/linear_fp8.py @@ -32,8 +32,8 @@ def fp8_matmul( 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_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - return fp8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale) + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) + return fp8_matmul(input, self.weight, self.bias, self.scale) fp8_matmul = compile_func(fp8_matmul) diff --git a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py index a35eb0721..d39ea5519 100644 --- a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py +++ b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py @@ -37,8 +37,8 @@ def fp8_matmul_tensorwise( def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor: if torch.numel(input) / input.shape[-1] < 32: - return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_dequantizer.scale) + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) + return fp8_matmul_tensorwise(input, self.weight, self.bias, self.scale) fp8_matmul_tensorwise = compile_func(fp8_matmul_tensorwise) diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py index 356179b97..bfca56ce2 100644 --- a/modules/sdnq/layers/linear/linear_int8.py +++ b/modules/sdnq/layers/linear/linear_int8.py @@ -41,13 +41,13 @@ def int8_matmul( def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: if torch.numel(input) / input.shape[-1] < 32: - return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) if self.sdnq_dequantizer.re_quantize_for_matmul: - weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight) + weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point) quantized_weight_shape = None else: weight = self.weight - scale = self.sdnq_dequantizer.scale + scale = self.scale quantized_weight_shape = getattr(self.sdnq_dequantizer, "quantized_weight_shape", None) return int8_matmul(input, weight, self.bias, scale, quantized_weight_shape, self.sdnq_dequantizer.weights_dtype)