diff --git a/modules/sdnq/quant_utils.py b/modules/sdnq/quant_utils.py index c702831f4..b0490ae4a 100644 --- a/modules/sdnq/quant_utils.py +++ b/modules/sdnq/quant_utils.py @@ -74,15 +74,14 @@ def apply_svdquant(weight: torch.FloatTensor, rank: int = 32, niter: int = 8, dt return weight, svd_up, svd_down -HADAMARD_N2_MATRIX = [[1, 1], [1, -1]] @devices.inference_context() def build_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None): if n == 1: return torch.ones((1, 1), dtype=dtype, device=device) - H = torch.tensor(HADAMARD_N2_MATRIX, dtype=dtype, device=device) current_size = 2 + H = H_N2 = torch.tensor([[1, 1], [1, -1]], dtype=dtype, device=device) while current_size < n: - H = torch.kron(H, torch.tensor(HADAMARD_N2_MATRIX, dtype=dtype, device=device)) + H = torch.kron(H, H_N2) current_size *= 2 H = H.div_(n**0.5) H = prepare_weight_for_matmul(H) @@ -97,13 +96,12 @@ HADAMARD_MATRIX_CACHE = {} def get_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None): global HADAMARD_MATRIX_CACHE # pylint: disable=global-variable-not-assigned device = devices.normalize_device(device) - if HADAMARD_MATRIX_CACHE.get(n, None) is None: - HADAMARD_MATRIX_CACHE[n] = {} - if HADAMARD_MATRIX_CACHE[n].get(device, None) is None: - HADAMARD_MATRIX_CACHE[n][device] = {} - if HADAMARD_MATRIX_CACHE[n][device].get(dtype, None) is None: - HADAMARD_MATRIX_CACHE[n][device][dtype] = build_hadamard(n, dtype=dtype, device=device) - return HADAMARD_MATRIX_CACHE[n][device][dtype] + H_key = (n, device, dtype) + H = HADAMARD_MATRIX_CACHE.get(H_key, None) + if H is None: + H = build_hadamard(n, dtype=dtype, device=device) + HADAMARD_MATRIX_CACHE[H_key] = H + return H @devices.inference_context() diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index 3392331a7..a9e9eec55 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -46,7 +46,27 @@ class QuantizationMethod(str, Enum): @devices.inference_context() -def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int8", quantized_matmul_dtype=None, torch_dtype=None, group_size=0, hadamard_group_size=128, svd_rank=32, svd_steps=8, use_svd=False, use_hadamard=False, use_quantized_matmul=False, use_stochastic_rounding=False, dequantize_fp32=True, using_pre_calculated_svd=False, using_pre_rotated_hadamard=False, skip_sr=False, param_name=None): # pylint: disable=unused-argument +def sdnq_quantize_layer_weight( + weight: torch.FloatTensor, + layer_class_name: str | None = None, + weights_dtype: str = "int8", + quantized_matmul_dtype: str | None = None, + group_size: int = 0, + hadamard_group_size: int = 128, + svd_rank: int = 32, + svd_steps: int = 8, + use_svd: bool = False, + use_hadamard: bool = False, + use_quantized_matmul: bool = False, + use_stochastic_rounding: bool = False, + dequantize_fp32: bool = True, + hadamard: torch.FloatTensor | None = None, + using_pre_calculated_svd: bool = False, + using_pre_rotated_hadamard: bool = False, + skip_sr: bool = False, + param_name: str | None = None, # pylint: disable=unused-argument + torch_dtype: torch.dtype | None = None, +): num_of_groups = 1 is_conv_type = False is_conv_transpose_type = False @@ -54,9 +74,9 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int result_shape = None scale_dtype = None + weight = weight.detach() original_shape = weight.shape original_stride = weight.stride() - weight = weight.detach() if torch_dtype is None: torch_dtype = weight.dtype @@ -122,7 +142,7 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int scale_dtype = torch_dtype if use_hadamard: - weight, use_hadamard, hadamard_group_size = apply_hadamard(weight, group_size=hadamard_group_size, layer_class_name=layer_class_name) + weight, use_hadamard, hadamard_group_size = apply_hadamard(weight, group_size=hadamard_group_size, hadamard=hadamard, layer_class_name=layer_class_name) if use_svd: try: @@ -228,20 +248,38 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int @devices.inference_context() -def sdnq_quantize_layer_weight_dynamic(weight, layer_class_name=None, weights_dtype="uint4", quantized_matmul_dtype=None, group_size=0, hadamard_group_size=128, svd_rank=32, svd_steps=8, dynamic_loss_threshold=None, use_svd=False, use_hadamard=False, use_quantized_matmul=False, use_stochastic_rounding=False, dequantize_fp32=True, quantization_config=None, torch_dtype=None, param_name=None): # pylint: disable=unused-argument +def sdnq_quantize_layer_weight_dynamic( + weight: torch.FloatTensor, + layer_class_name: str | None = None, + weights_dtype: str = "uint4", + quantized_matmul_dtype: str | None = None, + group_size: int = 0, + hadamard_group_size: int = 128, + svd_rank: int = 32, + svd_steps: int = 8, + dynamic_loss_threshold: float | None =None, + use_svd: bool = False, + use_hadamard: bool = False, + use_quantized_matmul: bool = False, + use_stochastic_rounding: bool = False, + dequantize_fp32: bool = True, + hadamard: torch.FloatTensor | None = None, + param_name: str | None = None, + torch_dtype: torch.dtype | None = None, +): if torch_dtype is None: torch_dtype = weight.dtype if dynamic_loss_threshold is None or dynamic_loss_threshold < 0: dynamic_loss_threshold = 10 ** -(dtype_dict[weights_dtype]["num_bits"] / 2) + weight = weight.detach() if weight.dtype != torch.float64: weight = weight.to(dtype=torch.float32, copy=False) - weight = weight.detach() - original_weight_fp32 = weight.clone() if use_svd else weight - weight_std = original_weight_fp32.std().square_().clamp_(min=1e-8) + weight_std = weight.std().square_().clamp_(min=1e-8) + original_weight_fp32 = weight if use_hadamard: - weight, use_hadamard, hadamard_group_size = apply_hadamard(weight, group_size=hadamard_group_size, layer_class_name=layer_class_name) + weight, use_hadamard, hadamard_group_size = apply_hadamard(weight, group_size=hadamard_group_size, hadamard=hadamard, layer_class_name=layer_class_name) if use_svd: try: @@ -575,10 +613,8 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): else: return_dtype = kwargs.get("dtype", param_value.dtype if self.torch_dtype is None else self.torch_dtype) - if param_value.dtype == return_dtype and devices.same_device(param_value.device, target_device): - param_value = param_value.detach() - else: - param_value = param_value.to(target_device, dtype=return_dtype, copy=False) + param_value = param_value.detach() + param_value = param_value.to(target_device, dtype=return_dtype, copy=False) if tensor_name == "weight" and layer.sdnq_dequantizer.use_quantized_matmul and not layer.sdnq_dequantizer.re_quantize_for_matmul: param_value = prepare_weight_for_matmul(param_value) @@ -600,10 +636,8 @@ class SDNQQuantizer(DiffusersQuantizer, HfQuantizer): target_device = quant_kwargs["quantization_device"] quant_kwargs["quantization_device"] = None - if param_value.dtype in {torch.float32, torch.float64} and devices.same_device(param_value.device, target_device): - param_value = param_value.detach() - else: - param_value = param_value.to(target_device, non_blocking=self.quantization_config.non_blocking, copy=False).to(dtype=torch.float32 if param_value.dtype != torch.float64 else torch.float64) + param_value = param_value.detach() + param_value = param_value.to(target_device, non_blocking=self.quantization_config.non_blocking, copy=False).to(dtype=torch.float32 if param_value.dtype != torch.float64 else torch.float64) layer.weight = torch.nn.Parameter(param_value, requires_grad=False) layer, self.quantization_config = sdnq_quantize_layer(layer, self.quantization_config, torch_dtype=torch_dtype, param_name=param_name, quant_kwargs=quant_kwargs) # pylint: disable=attribute-defined-outside-init