Update SDNQ

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

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

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):

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