Make SDNQ scales compatible with balanced offload

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

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

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 = {

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,

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,

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,

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)

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)

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)

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)

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)