diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 3afb60a67..6faca9887 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -5,7 +5,7 @@ import torch from modules import shared, devices -sdnq_version = "0.1.4" +sdnq_version = "0.1.5" dtype_dict = { ### Integers diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index b298ac1a2..ff1036260 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -9,6 +9,7 @@ from modules import devices from .common import dtype_dict, compile_func, use_contiguous_mm, use_tensorwise_fp8_matmul from .packed_int import unpack_int_symetric, unpack_int_asymetric from .packed_float import unpack_float +from .layers import SDNQLayer @devices.inference_context() @@ -95,7 +96,7 @@ def quantize_int_mm(input: torch.FloatTensor, dim: int = -1, matmul_dtype: str = @devices.inference_context() def quantize_int_mm_sr(input: torch.FloatTensor, dim: int = -1, matmul_dtype: str = "int8") -> Tuple[torch.Tensor, torch.FloatTensor]: scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(dtype_dict[matmul_dtype]["max"]) - input = torch.div(input, scale).add_(torch.rand_like(input), alpha=0.1).round_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"]) + input = torch.div(input, scale).add_(torch.randn_like(input), alpha=0.1).round_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"]) return input, scale @@ -111,7 +112,7 @@ def quantize_fp_mm_sr(input: torch.FloatTensor, dim: int = -1, matmul_dtype: str mantissa_difference = 1 << (23 - dtype_dict[matmul_dtype]["mantissa"]) scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(dtype_dict[matmul_dtype]["max"]) input = torch.div(input, scale).to(dtype=torch.float32).view(dtype=torch.int32) - input = input.add_(torch.randint_like(input, low=0, high=mantissa_difference, dtype=torch.int32)).view(dtype=torch.float32) + input = input.add_(torch.randint_like(input, low=0, high=mantissa_difference, dtype=torch.int32)).bitwise_and_(-mantissa_difference).view(dtype=torch.float32) input = input.nan_to_num_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"]) return input, scale @@ -177,30 +178,16 @@ def re_quantize_matmul_packed_float_symmetric(weight: torch.ByteTensor, scale: t return re_quantize_matmul_symmetric(unpack_float(weight, shape, weights_dtype), scale, matmul_dtype, svd_up=svd_up, svd_down=svd_down, result_shape=result_shape) -@devices.inference_context() -def dequantize_layer_weight(self: torch.nn.Module, inplace: bool = False): - weight = torch.nn.Parameter(self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul), requires_grad=True) - forward = getattr(torch.nn, self.sdnq_dequantizer.layer_class_name).forward - if inplace: - self.weight = weight - self.forward = forward - self.forward = self.forward.__get__(self, self.__class__) - del self.sdnq_dequantizer, self.scale, self.zero_point, self.svd_up, self.svd_down - return self - else: - return weight, forward - - @devices.inference_context() def dequantize_sdnq_module(model: torch.nn.Module): - if hasattr(model, "sdnq_dequantizer"): - model = dequantize_layer_weight(model, inplace=True) + if isinstance(model, SDNQLayer): + model = model.dequantize() has_children = list(model.children()) if not has_children: return model for module_name, module in model.named_children(): - if hasattr(module, "sdnq_dequantizer"): - setattr(model, module_name, dequantize_layer_weight(module, inplace=True)) + if isinstance(module, SDNQLayer): + setattr(model, module_name, module.dequantize()) else: setattr(model, module_name, dequantize_sdnq_model(module)) return model diff --git a/modules/sdnq/layers/__init__.py b/modules/sdnq/layers/__init__.py index 47f4420af..3e8ca9c76 100644 --- a/modules/sdnq/layers/__init__.py +++ b/modules/sdnq/layers/__init__.py @@ -1,3 +1,4 @@ +import copy import torch @@ -5,16 +6,26 @@ class SDNQLayer(torch.nn.Module): def __init__(self, original_layer, forward_func): torch.nn.Module.__init__(self) for key, value in original_layer.__dict__.items(): - if key not in {"forward", "forward_func", "original_class"}: + if key not in {"forward", "forward_func", "original_class", "state_dict", "load_state_dict"}: setattr(self, key, value) self.original_class = original_layer.__class__ self.forward_func = forward_func + def dequantize(self: torch.nn.Module): + if self.weight.__class__.__name__ == "SDNQTensor": + self.weight = torch.nn.Parameter(self.weight.dequantize(), requires_grad=True) + elif hasattr(self, "sdnq_dequantizer"): + self.weight = torch.nn.Parameter(self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul), requires_grad=True) + del self.sdnq_dequantizer, self.scale, self.zero_point, self.svd_up, self.svd_down + self.__class__ = self.original_class + del self.original_class, self.forward_func + return self + def forward(self, *args, **kwargs) -> torch.Tensor: return self.forward_func(self, *args, **kwargs) def __repr__(self): - return f"{self.__class__.__name__}(original_class={self.original_class.__name__} forward_func={self.forward_func} sdnq_dequantizer={repr(getattr(self, 'sdnq_dequantizer', None))})" + return f"{self.__class__.__name__}(original_class={self.original_class} forward_func={self.forward_func} sdnq_dequantizer={repr(getattr(self, 'sdnq_dequantizer', None))})" class SDNQLinear(SDNQLayer, torch.nn.Linear): diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index 9414ff6f6..95a644c01 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -56,12 +56,12 @@ def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[i if dtype_dict[weights_dtype]["is_integer"]: if use_stochastic_rounding: - quantized_weight.add_(torch.rand_like(quantized_weight), alpha=0.1) + quantized_weight.add_(torch.randn_like(quantized_weight), alpha=0.1) quantized_weight.round_() else: if use_stochastic_rounding: mantissa_difference = 1 << (23 - dtype_dict[weights_dtype]["mantissa"]) - quantized_weight = quantized_weight.view(dtype=torch.int32).add_(torch.randint_like(quantized_weight, low=0, high=mantissa_difference, dtype=torch.int32)).view(dtype=torch.float32) + quantized_weight = quantized_weight.view(dtype=torch.int32).add_(torch.randint_like(quantized_weight, low=0, high=mantissa_difference, dtype=torch.int32)).bitwise_and_(-mantissa_difference).view(dtype=torch.float32) quantized_weight.nan_to_num_() quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"]) return quantized_weight, scale, zero_point @@ -205,7 +205,7 @@ def add_module_skip_keys(model, modules_to_not_convert: List[str] = None, module @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, svd_rank=32, svd_steps=8, use_svd=False, use_quantized_matmul=False, use_stochastic_rounding=False, dequantize_fp32=False, using_pre_calculated_svd=False, param_name=None): # pylint: disable=unused-argument +def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int8", quantized_matmul_dtype=None, torch_dtype=None, group_size=0, svd_rank=32, svd_steps=8, use_svd=False, use_quantized_matmul=False, use_stochastic_rounding=False, dequantize_fp32=False, using_pre_calculated_svd=False, skip_sr=False, param_name=None): # pylint: disable=unused-argument num_of_groups = 1 is_conv_type = False is_conv_transpose_type = False @@ -335,7 +335,7 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int else: group_size = -1 - weight, scale, zero_point = quantize_weight(weight, reduction_axes, weights_dtype, use_stochastic_rounding=use_stochastic_rounding) + weight, scale, zero_point = quantize_weight(weight, reduction_axes, weights_dtype, use_stochastic_rounding=(use_stochastic_rounding and not skip_sr)) if ( not dequantize_fp32 and dtype_dict[weights_dtype]["num_bits"] <= 8