From 9e52d0c1fbf76f2f9a2cccc224f0cd3cb71703c4 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 5 Oct 2025 22:50:30 +0300 Subject: [PATCH] SDNQ add SVDQuant quantization method --- CHANGELOG.md | 5 +- modules/lora/lora_apply.py | 53 +++++++++++-- modules/model_quant.py | 10 ++- modules/sdnq/__init__.py | 75 ++++++++++++++---- modules/sdnq/dequantizer.py | 78 +++++++++++-------- modules/sdnq/layers/conv/conv_fp8.py | 12 ++- .../sdnq/layers/conv/conv_fp8_tensorwise.py | 13 +++- modules/sdnq/layers/conv/conv_int8.py | 15 +++- modules/sdnq/layers/conv/forward.py | 8 +- modules/sdnq/layers/linear/forward.py | 2 +- modules/sdnq/layers/linear/linear_fp8.py | 13 +++- .../layers/linear/linear_fp8_tensorwise.py | 11 ++- modules/sdnq/layers/linear/linear_int8.py | 13 +++- modules/shared.py | 2 + 14 files changed, 229 insertions(+), 81 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4d3d9a478..55df207f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -59,9 +59,10 @@ - separate guidance and detail sections - networks ability to filter lora by base model version - **SDNQ** + - add `SVDQuant` quantization method support - make sdnq scales compatible with balanced offload - - add int8 matmul support for rdna2 gpus via triton - - improve int8 mamtul performance on intel gpus + - add int8 matmul support for RDNA2 GPUs via triton + - improve int8 mamtul performance on Intel GPUs - **Other** - server will note when restart is recommended due to package updates - **interrrupt** will now show last known preview image diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index 47b408446..d076efb97 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -52,6 +52,12 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n self.sdnq_zero_point_backup = self.zero_point.clone().to(devices.cpu) else: self.sdnq_zero_point_backup = None + if self.svd_up is not None: + self.sdnq_svd_up_backup = self.svd_up.clone().to(devices.cpu) + self.sdnq_svd_down_backup = self.svd_down.clone().to(devices.cpu) + else: + self.sdnq_svd_up_backup = None + self.sdnq_svd_down_backup = None if bias_backup is None: if getattr(self, 'bias', None) is not None: @@ -85,9 +91,23 @@ 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), 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) + 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, + self.sdnq_svd_up_backup.to(devices.device) if self.sdnq_svd_up_backup is not None else None, + self.sdnq_svd_down_backup.to(devices.device) if self.sdnq_svd_down_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), 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) + 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, + self.svd_up.to(devices.device) if self.svd_up is not None else None, + self.svd_down.to(devices.device) if self.svd_down 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) @@ -148,19 +168,35 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G from modules.sdnq import sdnq_quantize_layer if hasattr(self, "sdnq_dequantizer_backup"): 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) + 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, + self.sdnq_svd_up_backup.to(devices.device) if self.sdnq_svd_up_backup is not None else None, + self.sdnq_svd_down_backup.to(devices.device) if self.sdnq_svd_down_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) + 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, + self.svd_up.to(devices.device) if self.svd_up is not None else None, + self.svd_down.to(devices.device) if self.svd_down 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) - del self.sdnq_dequantizer, self.scale, self.zero_point + del self.sdnq_dequantizer, self.scale, self.zero_point, self.svd_up, self.svd_down self = sdnq_quantize_layer( self, weights_dtype=weights_dtype, torch_dtype=devices.dtype, group_size=shared.opts.sdnq_quantize_weights_group_size, + svd_rank=shared.opts.sdnq_svd_rank, + use_svd=shared.opts.sdnq_use_svd, quant_conv=shared.opts.sdnq_quantize_conv_layers, use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv, @@ -243,7 +279,12 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn 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 self.sdnq_svd_up_backup is not None: + self.svd_up = torch.nn.Parameter(self.sdnq_svd_up_backup.to(device), requires_grad=False) + self.svd_down = torch.nn.Parameter(self.sdnq_svd_down_backup.to(device), requires_grad=False) + else: + self.svd_up, self.svd_down = None, None + del self.sdnq_dequantizer_backup, self.sdnq_scale_backup, self.sdnq_zero_point_backup, self.sdnq_svd_up_backup, self.sdnq_svd_down_backup if bias_backup is not None: self.bias = None diff --git a/modules/model_quant.py b/modules/model_quant.py index 19b7ae7a7..288d6e8ed 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -205,6 +205,8 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', sdnq_config = SDNQConfig( weights_dtype=weights_dtype, group_size=shared.opts.sdnq_quantize_weights_group_size, + svd_rank=shared.opts.sdnq_svd_rank, + use_svd=shared.opts.sdnq_use_svd, quant_conv=shared.opts.sdnq_quantize_conv_layers, use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv, @@ -215,7 +217,7 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict.copy(), ) - log.debug(f'Quantization: module="{module}" type=sdnq mode=pre dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device} device_map={shared.opts.device_map} offload_mode={shared.opts.diffusers_offload_mode} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_to_not_convert={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}') + log.debug(f'Quantization: module="{module}" type=sdnq mode=pre dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} svd_rank={shared.opts.sdnq_svd_rank} use_svd={shared.opts.sdnq_use_svd} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device} device_map={shared.opts.device_map} offload_mode={shared.opts.diffusers_offload_mode} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_to_not_convert={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}') if kwargs is None: return sdnq_config else: @@ -533,6 +535,8 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh weights_dtype=weights_dtype, torch_dtype=devices.dtype, group_size=shared.opts.sdnq_quantize_weights_group_size, + svd_rank=shared.opts.sdnq_svd_rank, + use_svd=shared.opts.sdnq_use_svd, quant_conv=shared.opts.sdnq_quantize_conv_layers, use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv, @@ -572,7 +576,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh if do_gc: devices.torch_gc(force=True, reason='sdnq') - log.debug(f'Quantization: module="{op if op is not None else model.__class__}" type=sdnq mode=post dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device} device_map={shared.opts.device_map} offload_mode={shared.opts.diffusers_offload_mode} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_to_not_convert={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}') + log.debug(f'Quantization: module="{op if op is not None else model.__class__}" type=sdnq mode=post dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} svd_rank={shared.opts.sdnq_svd_rank} use_svd={shared.opts.sdnq_use_svd} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device} device_map={shared.opts.device_map} offload_mode={shared.opts.diffusers_offload_mode} non_blocking={shared.opts.diffusers_offload_nonblocking} modules_to_not_convert={modules_to_not_convert} modules_dtype_dict={modules_dtype_dict}') return model @@ -580,7 +584,7 @@ def sdnq_quantize_weights(sd_model): try: t0 = time.time() from modules import shared, devices, sd_models - log.debug(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights} dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} pre_forward={shared.opts.diffusers_offload_pre}") + log.debug(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights} dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul} svd_rank={shared.opts.sdnq_svd_rank} use_svd={shared.opts.sdnq_use_svd} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} pre_forward={shared.opts.diffusers_offload_pre}") global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement sd_model = sd_models.apply_function_to_model(sd_model, sdnq_quantize_model, shared.opts.sdnq_quantize_weights, op="sdnq") diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index 39077d917..5e90e0cc1 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -48,8 +48,15 @@ def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[i return quantized_weight, scale, zero_point +def apply_svdquant(weight: torch.FloatTensor, rank: int = 128) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + U, S, svd_down = torch.svd_lowrank(weight, q=rank) + svd_up = torch.mul(U, S.unsqueeze(0)) + svd_down = svd_down.t_() + return weight.sub_(torch.mm(svd_up, svd_down)), svd_up, svd_down + + @devices.inference_context() -def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument +def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=128, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument layer_class_name = layer.__class__.__name__ if layer_class_name in allowed_types: num_of_groups = 1 @@ -103,13 +110,30 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz else: use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0 + layer.weight.requires_grad = False + if return_device is None: + return_device = layer.weight.device + if quantization_device is not None: + layer.weight.data = layer.weight.to(quantization_device, non_blocking=non_blocking) + if layer.weight.dtype != torch.float32: + layer.weight.data = layer.weight.to(dtype=torch.float32) + + if use_svd and is_linear_type: + layer.weight.data, svd_up, svd_down = apply_svdquant(layer.weight, rank=svd_rank) + if use_quantized_matmul: + svd_up = svd_up.t_() + svd_down = svd_down.t_() + else: + use_svd = False + svd_up, svd_down = None, None + if group_size == 0: if use_quantized_matmul and dtype_dict[weights_dtype]["num_bits"] >= 6: group_size = -1 elif is_linear_type: - group_size = 2 ** (2 + dtype_dict[weights_dtype]["num_bits"]) + group_size = 2 ** ((2 if not use_svd else 3) + dtype_dict[weights_dtype]["num_bits"]) else: - group_size = 2 ** (1 + dtype_dict[weights_dtype]["num_bits"]) + group_size = 2 ** ((1 if not use_svd else 2) + dtype_dict[weights_dtype]["num_bits"]) elif use_quantized_matmul and dtype_dict[weights_dtype]["num_bits"] == 8: group_size = -1 # override user value, re-quantizing 8bit into 8bit is pointless elif group_size != -1 and not is_linear_type: @@ -154,24 +178,19 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz new_shape[last_dim_index - 1 : last_dim_index] = (num_of_groups, group_size) layer.weight.data = layer.weight.reshape(new_shape) - layer.weight.requires_grad = False - if return_device is None: - return_device = layer.weight.device - if quantization_device is not None: - layer.weight.data = layer.weight.to(quantization_device, non_blocking=non_blocking) - if layer.weight.dtype != torch.float32: - layer.weight.data = layer.weight.to(dtype=torch.float32) - layer.weight.data, scale, zero_point = quantize_weight(layer.weight, reduction_axes, weights_dtype) if not dequantize_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"] and not use_tensorwise_fp8_matmul): - scale = scale.to(torch_dtype) + scale = scale.to(dtype=torch_dtype) if zero_point is not None: - zero_point = zero_point.to(torch_dtype) + zero_point = zero_point.to(dtype=torch_dtype) + if svd_up is not None: + svd_up = svd_up.to(dtype=torch_dtype) + svd_down = svd_down.to(dtype=torch_dtype) re_quantize_for_matmul = (num_of_groups > 1 or zero_point is not None) if use_quantized_matmul and not re_quantize_for_matmul: - scale.transpose_(0,1) - layer.weight.transpose_(0,1) + scale.t_() + layer.weight.t_() if use_contiguous_mm: layer.weight.data = layer.weight.contiguous() elif layer.weight.is_contiguous(): @@ -186,6 +205,13 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz layer.zero_point = torch.nn.Parameter(zero_point, requires_grad=False) else: layer.zero_point = None + if svd_up is not None: + svd_up = svd_up.to(return_device, non_blocking=non_blocking) + svd_down = svd_down.to(return_device, non_blocking=non_blocking) + layer.svd_up = torch.nn.Parameter(svd_up, requires_grad=False) + layer.svd_down = torch.nn.Parameter(svd_down, requires_grad=False) + else: + layer.svd_up, layer.svd_down = None, None layer.sdnq_dequantizer = dequantizer_dict[weights_dtype]( quantized_weight_shape=layer.weight.shape, @@ -204,7 +230,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz return layer -def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, modules_to_not_convert: List[str] = None, modules_dtype_dict: Dict[str, List[str]] = None, op=None): # pylint: disable=unused-argument +def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, svd_rank=128, use_svd=False, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, modules_to_not_convert: List[str] = None, modules_dtype_dict: Dict[str, List[str]] = None, op=None): # pylint: disable=unused-argument has_children = list(model.children()) if not has_children: return model @@ -238,6 +264,8 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si weights_dtype=weights_dtype, torch_dtype=torch_dtype, group_size=group_size, + svd_rank=svd_rank, + use_svd=use_svd, quant_conv=quant_conv, use_quantized_matmul=use_quantized_matmul, use_quantized_matmul_conv=use_quantized_matmul_conv, @@ -252,6 +280,8 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si weights_dtype=weights_dtype, torch_dtype=torch_dtype, group_size=group_size, + svd_rank=svd_rank, + use_svd=use_svd, quant_conv=quant_conv, use_quantized_matmul=use_quantized_matmul, use_quantized_matmul_conv=use_quantized_matmul_conv, @@ -371,6 +401,8 @@ class SDNQQuantizer(DiffusersQuantizer): weights_dtype=weights_dtype, torch_dtype=self.torch_dtype, group_size=self.quantization_config.group_size, + svd_rank=self.quantization_config.svd_rank, + use_svd=self.quantization_config.use_svd, quant_conv=self.quantization_config.quant_conv, use_quantized_matmul=self.quantization_config.use_quantized_matmul, use_quantized_matmul_conv=self.quantization_config.use_quantized_matmul_conv, @@ -497,8 +529,13 @@ class SDNQConfig(QuantizationConfigMixin): weights_dtype (`str`, *optional*, defaults to `"int8"`): The target dtype for the weights after quantization. Supported values are: ("int8", "int7", "int6", "int5", "int4", "int3", "int2", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz") - weights_dtype (`int`, *optional*, defaults to `0`): + group_size (`int`, *optional*, defaults to `0`): Used to decide how many elements of a tensor will share the same quantization group. + group_size = 0 will automatically select a group size based on weights_dtype. + svd_rank (`int`, *optional*, defaults to `128`): + The rank size used for the SVDQuant algorithm. + use_svd (`bool`, *optional*, defaults to `False`): + Enabling this option will use SVDQuant algorithm. quant_conv (`bool`, *optional*, defaults to `False`): Enabling this option will quantize the convolutional layers in UNet models too. use_quantized_matmul (`bool`, *optional*, defaults to `False`): @@ -524,6 +561,8 @@ class SDNQConfig(QuantizationConfigMixin): self, weights_dtype: str = "int8", group_size: int = 0, + svd_rank: int = 128, + use_svd: bool = False, quant_conv: bool = False, use_quantized_matmul: bool = False, use_quantized_matmul_conv: bool = False, @@ -538,6 +577,8 @@ class SDNQConfig(QuantizationConfigMixin): self.weights_dtype = weights_dtype self.quant_method = QuantizationMethod.SDNQ self.group_size = group_size + self.svd_rank = svd_rank + self.use_svd = use_svd self.quant_conv = quant_conv self.use_quantized_matmul = use_quantized_matmul self.use_quantized_matmul_conv = use_quantized_matmul_conv diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index 3881cfe34..72a058555 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -1,6 +1,6 @@ # pylint: disable=redefined-builtin,no-member,protected-access -from typing import Tuple +from typing import Tuple, Optional import torch @@ -8,19 +8,31 @@ from .common import dtype_dict, compile_func from .packed_int import pack_int_symetric, unpack_int_symetric, pack_int_asymetric, unpack_int_asymetric -def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor: - result = torch.addcmul(zero_point, weight.to(dtype=scale.dtype), scale).to(dtype=dtype) +def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor: + result = torch.addcmul(zero_point, weight.to(dtype=scale.dtype), scale) if result_shape is not None: result = result.view(result_shape) + if svd_up is not None: + if skip_quantized_matmul: + svd_up, svd_down = svd_up.t(), svd_down.t() + result = torch.addmm(result, svd_up, svd_down) + if dtype is not None: + result = result.to(dtype=dtype) return result -def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor: - result = weight.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype) +def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor: + result = weight.to(dtype=scale.dtype).mul_(scale) if skip_quantized_matmul: result.t_() if result_shape is not None: result = result.view(result_shape) + if svd_up is not None: + if skip_quantized_matmul: + svd_up, svd_down = svd_up.t(), svd_down.t() + result = torch.addmm(result, svd_up, svd_down) + if dtype is not None: + result = result.to(dtype=dtype) return result @@ -28,12 +40,12 @@ def dequantize_symmetric_with_bias(weight: torch.CharTensor, scale: torch.FloatT return torch.addcmul(bias, weight.to(dtype=scale.dtype), scale).to(dtype=dtype).view(result_shape) -def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor: - return dequantize_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, dtype, result_shape) +def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor: + return dequantize_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, dtype, result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul) -def dequantize_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.FloatTensor: - return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape, skip_quantized_matmul=skip_quantized_matmul) +def dequantize_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None, skip_quantized_matmul: bool = False) -> torch.FloatTensor: + return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul) def quantize_int8(input: torch.FloatTensor, dim: int = -1) -> Tuple[torch.CharTensor, torch.FloatTensor]: @@ -55,20 +67,20 @@ def re_quantize_int8(weight: torch.FloatTensor) -> Tuple[torch.CharTensor, torch return weight, scale -def re_quantize_matmul_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, result_shape: torch.Size) -> Tuple[torch.CharTensor, torch.FloatTensor]: - return re_quantize_int8(dequantize_asymmetric(weight, scale, zero_point, scale.dtype, result_shape)) +def re_quantize_matmul_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]: + return re_quantize_int8(dequantize_asymmetric(weight, scale, zero_point, scale.dtype, result_shape, svd_up=svd_up, svd_down=svd_down)) -def re_quantize_matmul_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, result_shape: torch.Size) -> Tuple[torch.CharTensor, torch.FloatTensor]: - return re_quantize_int8(dequantize_symmetric(weight, scale, scale.dtype, result_shape)) +def re_quantize_matmul_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, result_shape: torch.Size, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> Tuple[torch.CharTensor, torch.FloatTensor]: + return re_quantize_int8(dequantize_symmetric(weight, scale, scale.dtype, result_shape, svd_up=svd_up, svd_down=svd_down)) -def re_quantize_matmul_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor: - return re_quantize_matmul_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, result_shape) +def re_quantize_matmul_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> torch.FloatTensor: + return re_quantize_matmul_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, result_shape, svd_up=svd_up, svd_down=svd_down) -def re_quantize_matmul_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor: - return re_quantize_matmul_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, result_shape) +def re_quantize_matmul_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, result_shape: torch.Size, weights_dtype: str, svd_up: Optional[torch.FloatTensor] = None, svd_down: Optional[torch.FloatTensor] = None) -> torch.FloatTensor: + return re_quantize_matmul_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, result_shape, svd_up=svd_up, svd_down=svd_down) class AsymmetricWeightsDequantizer(torch.nn.Module): @@ -92,11 +104,11 @@ class AsymmetricWeightsDequantizer(torch.nn.Module): 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, scale, zero_point, **kwargs): # pylint: disable=unused-argument - return re_quantize_matmul_asymmetric_compiled(weight, scale, zero_point, self.result_shape) + def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **kwargs): # pylint: disable=unused-argument + return re_quantize_matmul_asymmetric_compiled(weight, scale, zero_point, self.result_shape, svd_up=svd_up, svd_down=svd_down) - 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) + def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument + return dequantize_asymmetric_compiled(weight, scale, zero_point, self.result_dtype, self.result_shape, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul) class SymmetricWeightsDequantizer(torch.nn.Module): @@ -121,12 +133,12 @@ class SymmetricWeightsDequantizer(torch.nn.Module): 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, scale, zero_point, **kwargs): # pylint: disable=unused-argument - return re_quantize_matmul_symmetric_compiled(weight, scale, self.result_shape) + def re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **kwargs): # pylint: disable=unused-argument + return re_quantize_matmul_symmetric_compiled(weight, scale, self.result_shape, svd_up=svd_up, svd_down=svd_down) - def forward(self, weight, scale, zero_point, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument + def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument skip_quantized_matmul = skip_quantized_matmul and not self.re_quantize_for_matmul - return dequantize_symmetric_compiled(weight, 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, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul) class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module): @@ -152,11 +164,11 @@ class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module): def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: return pack_int_asymetric(weight, 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 re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **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, svd_up=svd_up, svd_down=svd_down) - 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) + def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # 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, svd_up=svd_up, svd_down=svd_down, skip_quantized_matmul=skip_quantized_matmul) class PackedINTSymmetricWeightsDequantizer(torch.nn.Module): @@ -183,12 +195,12 @@ class PackedINTSymmetricWeightsDequantizer(torch.nn.Module): def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: return pack_int_symetric(weight, 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 re_quantize_matmul(self, weight, scale, zero_point, svd_up, svd_down, **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, svd_up=svd_up, svd_down=svd_down) - def forward(self, weight, scale, zero_point, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument + def forward(self, weight, scale, zero_point, svd_up, svd_down, skip_quantized_matmul=False): # pylint: disable=unused-argument skip_quantized_matmul = skip_quantized_matmul and not self.re_quantize_for_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) + return dequantize_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, 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 47f214937..d700a6998 100644 --- a/modules/sdnq/layers/conv/conv_fp8.py +++ b/modules/sdnq/layers/conv/conv_fp8.py @@ -15,6 +15,8 @@ def conv_fp8_matmul( weight: torch.Tensor, bias: torch.FloatTensor, scale: torch.FloatTensor, + svd_up: torch.FloatTensor, + svd_down: torch.FloatTensor, result_shape: torch.Size, reversed_padding_repeated_twice: List[int], padding_mode: str, conv_type: int, @@ -23,6 +25,9 @@ def conv_fp8_matmul( ) -> torch.FloatTensor: return_dtype = input.dtype input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) + if svd_up is not None: + svd_bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up) + input, input_scale = quantize_fp8_matmul_input(input) input, weight = check_mats(input, weight) @@ -46,6 +51,8 @@ def conv_fp8_matmul( for i in range(groups): result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=torch.bfloat16)) result = torch.cat(result, dim=-1).view(mm_output_shape).to(return_dtype) + if svd_up is not None: + result.add_(svd_bias) if conv_type == 1: result = result.transpose_(1,2) @@ -58,10 +65,11 @@ 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, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, 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.scale, + input, self.weight, self.bias, + self.scale, self.svd_up, self.svd_down, 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 592a512ee..1893ad8a5 100644 --- a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py +++ b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py @@ -16,6 +16,8 @@ def conv_fp8_matmul_tensorwise( weight: torch.Tensor, bias: torch.FloatTensor, scale: torch.FloatTensor, + svd_up: torch.FloatTensor, + svd_down: torch.FloatTensor, result_shape: torch.Size, reversed_padding_repeated_twice: List[int], padding_mode: str, conv_type: int, @@ -24,6 +26,12 @@ def conv_fp8_matmul_tensorwise( ) -> torch.FloatTensor: return_dtype = input.dtype input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) + if svd_up is not None: + if bias is not None: + bias = torch.addmm(bias, torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up) + else: + bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up) + input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) input, weight = check_mats(input, weight) dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) @@ -53,10 +61,11 @@ 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, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, 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.scale, + input, self.weight, self.bias, + self.scale, self.svd_up, self.svd_down, 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 294f696b4..82bd51f95 100644 --- a/modules/sdnq/layers/conv/conv_int8.py +++ b/modules/sdnq/layers/conv/conv_int8.py @@ -17,6 +17,8 @@ def conv_int8_matmul( weight: torch.CharTensor, bias: torch.FloatTensor, scale: torch.FloatTensor, + svd_up: torch.FloatTensor, + svd_down: torch.FloatTensor, quantized_weight_shape: torch.Size, result_shape: torch.Size, weights_dtype: str, @@ -27,6 +29,12 @@ def conv_int8_matmul( ) -> torch.FloatTensor: return_dtype = input.dtype input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) + if svd_up is not None: + if bias is not None: + bias = torch.addmm(bias, torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up) + else: + bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up) + input, scale = quantize_int8_matmul_input(input, scale) if quantized_weight_shape is not None: weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8) @@ -57,10 +65,10 @@ 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, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, 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, self.scale, self.zero_point) + weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None) quantized_weight_shape = None else: weight = self.weight @@ -68,7 +76,8 @@ def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor: quantized_weight_shape = getattr(self.sdnq_dequantizer, "quantized_weight_shape", None) return conv_int8_matmul( input, weight, self.bias, - scale, quantized_weight_shape, + scale, self.svd_up, self.svd_down, + quantized_weight_shape, self.sdnq_dequantizer.result_shape, self.sdnq_dequantizer.weights_dtype, self._reversed_padding_repeated_twice, diff --git a/modules/sdnq/layers/conv/forward.py b/modules/sdnq/layers/conv/forward.py index 69237ea68..2ed3d816f 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.scale, self.zero_point), self.bias) + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), 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.scale, self.zero_point), 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.svd_up, self.svd_down), 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.scale, self.zero_point), 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.svd_up, self.svd_down), 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.scale, self.zero_point), 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.svd_up, self.svd_down), 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 880b8ebc8..2371b6abb 100644 --- a/modules/sdnq/layers/linear/forward.py +++ b/modules/sdnq/layers/linear/forward.py @@ -17,4 +17,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.scale, self.zero_point), self.bias) + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down), self.bias) diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py index 3f2fb407c..205b141e8 100644 --- a/modules/sdnq/layers/linear/linear_fp8.py +++ b/modules/sdnq/layers/linear/linear_fp8.py @@ -20,20 +20,27 @@ def fp8_matmul( weight: torch.Tensor, bias: torch.FloatTensor, scale: torch.FloatTensor, + svd_up: torch.FloatTensor, + svd_down: torch.FloatTensor, ) -> torch.FloatTensor: return_dtype = input.dtype output_shape = (*input.shape[:-1], weight.shape[-1]) + if svd_up is not None: + svd_bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up) input, input_scale = quantize_fp8_matmul_input(input) input, weight = check_mats(input, weight) if bias is not None and bias.dtype != torch.bfloat16: bias = bias.to(dtype=torch.bfloat16) - return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16).view(output_shape).to(return_dtype) + result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16).view(output_shape).to(return_dtype) + if svd_up is not None: + result.add_(svd_bias) + return result 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, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) - return fp8_matmul(input, self.weight, self.bias, self.scale) + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias) + return fp8_matmul(input, self.weight, self.bias, self.scale, self.svd_up, self.svd_down) 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 d39ea5519..bd7a9b973 100644 --- a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py +++ b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py @@ -23,10 +23,17 @@ def fp8_matmul_tensorwise( weight: torch.Tensor, bias: torch.FloatTensor, scale: torch.FloatTensor, + svd_up: torch.FloatTensor, + svd_down: torch.FloatTensor, ) -> torch.FloatTensor: return_dtype = input.dtype output_shape = (*input.shape[:-1], weight.shape[-1]) dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) + if svd_up is not None: + if bias is not None: + bias = torch.addmm(bias, torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up) + else: + bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up) input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) input, weight = check_mats(input, weight) if bias is not None: @@ -37,8 +44,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, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) - return fp8_matmul_tensorwise(input, self.weight, self.bias, self.scale) + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias) + return fp8_matmul_tensorwise(input, self.weight, self.bias, self.scale, self.svd_up, self.svd_down) 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 ae218a4fe..ce8bcf788 100644 --- a/modules/sdnq/layers/linear/linear_int8.py +++ b/modules/sdnq/layers/linear/linear_int8.py @@ -24,6 +24,8 @@ def int8_matmul( weight: torch.Tensor, bias: torch.FloatTensor, scale: torch.FloatTensor, + svd_up: torch.FloatTensor, + svd_down: torch.FloatTensor, quantized_weight_shape: torch.Size, weights_dtype: str, ) -> torch.FloatTensor: @@ -31,6 +33,11 @@ def int8_matmul( weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8) return_dtype = input.dtype output_shape = (*input.shape[:-1], weight.shape[-1]) + if svd_up is not None: + if bias is not None: + bias = torch.addmm(bias, torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up) + else: + bias = torch.mm(torch.mm(input.flatten(0,-2).to(dtype=svd_down.dtype), svd_down), svd_up) input, scale = quantize_int8_matmul_input(input, scale) input, weight = check_mats(input, weight) if bias is not None: @@ -41,15 +48,15 @@ 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, self.scale, self.zero_point, skip_quantized_matmul=True), self.bias) + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias) if self.sdnq_dequantizer.re_quantize_for_matmul: - weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point) + weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None) quantized_weight_shape = None else: weight = self.weight 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) + return int8_matmul(input, weight, self.bias, scale, self.svd_up, self.svd_down, quantized_weight_shape, self.sdnq_dequantizer.weights_dtype) int8_matmul = compile_func(int8_matmul) diff --git a/modules/shared.py b/modules/shared.py index fab234f94..339cac0f1 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -201,6 +201,8 @@ options_templates.update(options_section(("quantization", "Model Quantization"), "sdnq_modules_to_not_convert": OptionInfo("", "Modules to not convert"), "sdnq_modules_dtype_dict": OptionInfo("{}", "Modules dtype dict"), "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1}), + "sdnq_svd_rank": OptionInfo(128, "SVDQuant Rank size", gr.Slider, {"minimum": 1, "maximum": 512, "step": 1}), + "sdnq_use_svd": OptionInfo(False, "Use SVDQuant quantization", gr.Checkbox), "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize convolutional layers", gr.Checkbox), "sdnq_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox), "sdnq_use_quantized_matmul": OptionInfo(False, "Use quantized MatMul", gr.Checkbox),