update sdnq

modules/sdnq/layers/__init__.py

@@ -0,0 +1,69 @@
+import torch
+
+
+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"}:
+                setattr(self, key, value)
+        self.original_class = original_layer.__class__
+        self.forward_func = forward_func
+
+    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))})"
+
+
+class SDNQLinear(SDNQLayer, torch.nn.Linear):
+    original_class: torch.nn.Linear
+
+class SDNQConv1d(SDNQLayer, torch.nn.Conv1d):
+    original_class: torch.nn.Conv1d
+
+class SDNQConv2d(SDNQLayer, torch.nn.Conv2d):
+    original_class: torch.nn.Conv2d
+
+class SDNQConv3d(SDNQLayer, torch.nn.Conv3d):
+    original_class: torch.nn.Conv3d
+
+class SDNQConvTranspose1d(SDNQLayer, torch.nn.ConvTranspose1d):
+    original_class: torch.nn.ConvTranspose1d
+
+class SDNQConvTranspose2d(SDNQLayer, torch.nn.ConvTranspose2d):
+    original_class: torch.nn.ConvTranspose2d
+
+class SDNQConvTranspose3d(SDNQLayer, torch.nn.ConvTranspose3d):
+    original_class: torch.nn.ConvTranspose3d
+
+
+torch.serialization.add_safe_globals([SDNQLayer])
+torch.serialization.add_safe_globals([SDNQLinear])
+torch.serialization.add_safe_globals([SDNQConv1d])
+torch.serialization.add_safe_globals([SDNQConv2d])
+torch.serialization.add_safe_globals([SDNQConv3d])
+torch.serialization.add_safe_globals([SDNQConvTranspose1d])
+torch.serialization.add_safe_globals([SDNQConvTranspose2d])
+torch.serialization.add_safe_globals([SDNQConvTranspose3d])
+
+
+def get_sdnq_wrapper_class(original_layer, forward_func):
+    match original_layer.__class__.__name__:
+        case "Linear":
+            return SDNQLinear(original_layer, forward_func)
+        case "Conv1d":
+            return SDNQConv1d(original_layer, forward_func)
+        case "Conv2d":
+            return SDNQConv2d(original_layer, forward_func)
+        case "Conv3d":
+            return SDNQConv3d(original_layer, forward_func)
+        case "ConvTranspose1d":
+            return SDNQConvTranspose1d(original_layer, forward_func)
+        case "ConvTranspose2d":
+            return SDNQConvTranspose2d(original_layer, forward_func)
+        case "ConvTranspose3d":
+            return SDNQConvTranspose3d(original_layer, forward_func)
+        case _:
+            return SDNQLayer(original_layer, forward_func)

modules/sdnq/layers/conv/conv_fp16.py

@@ -37,11 +37,12 @@ def conv_fp16_matmul(
         else:
             bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
 
-    input, scale = quantize_fp_mm_input_tensorwise(input, scale, matmul_dtype="float16")
     if quantized_weight_shape is not None:
-        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float16)
+        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float16).t_()
+        scale = scale.t()
     elif weight.dtype != torch.float16:
         weight = weight.to(dtype=torch.float16) # fp8 weights
+    input, scale = quantize_fp_mm_input_tensorwise(input, scale, matmul_dtype="float16")
     input, weight = check_mats(input, weight)
 
     if groups == 1:

modules/sdnq/layers/conv/conv_fp8.py

@@ -33,9 +33,10 @@ def conv_fp8_matmul(
         input = input.flatten(0,-2)
         svd_bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
 
-    input, input_scale = quantize_fp_mm_input(input)
     if quantized_weight_shape is not None:
-        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn)
+        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn).t_()
+        scale = scale.t()
+    input, input_scale = quantize_fp_mm_input(input)
     input, weight = check_mats(input, weight)
 
     if groups == 1:

modules/sdnq/layers/conv/conv_fp8_tensorwise.py

@@ -37,9 +37,10 @@ def conv_fp8_matmul_tensorwise(
         else:
             bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
 
-    input, scale = quantize_fp_mm_input_tensorwise(input, scale)
     if quantized_weight_shape is not None:
-        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn)
+        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn).t_()
+        scale = scale.t()
+    input, scale = quantize_fp_mm_input_tensorwise(input, scale)
     input, weight = check_mats(input, weight)
     dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
 

modules/sdnq/layers/conv/conv_int8.py

@@ -37,9 +37,10 @@ def conv_int8_matmul(
         else:
             bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
 
-    input, scale = quantize_int_mm_input(input, scale)
     if quantized_weight_shape is not None:
-        weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8)
+        weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8).t_()
+        scale = scale.t()
+    input, scale = quantize_int_mm_input(input, scale)
     input, weight = check_mats(input, weight)
 
     if groups == 1:

modules/sdnq/layers/linear/linear_fp16.py

@@ -21,7 +21,8 @@ def fp16_matmul(
     weights_dtype: str = None,
 ) -> torch.FloatTensor:
     if quantized_weight_shape is not None:
-        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float16)
+        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float16).t_()
+        scale = scale.t()
     elif weight.dtype != torch.float16:
         weight = weight.to(dtype=torch.float16) # fp8 weights
     return_dtype = input.dtype

modules/sdnq/layers/linear/linear_fp8.py

@@ -28,7 +28,8 @@ def fp8_matmul(
     weights_dtype: str = None,
 ) -> torch.FloatTensor:
     if quantized_weight_shape is not None:
-        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn)
+        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn).t_()
+        scale = scale.t()
     return_dtype = input.dtype
     output_shape = (*input.shape[:-1], weight.shape[-1])
     if svd_up is not None:

modules/sdnq/layers/linear/linear_fp8_tensorwise.py

@@ -31,7 +31,8 @@ def fp8_matmul_tensorwise(
     weights_dtype: str = None,
 ) -> torch.FloatTensor:
     if quantized_weight_shape is not None:
-        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn)
+        weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float8_e4m3fn).t_()
+        scale = scale.t()
     return_dtype = input.dtype
     output_shape = (*input.shape[:-1], weight.shape[-1])
     if svd_up is not None:

modules/sdnq/layers/linear/linear_int8.py

@@ -31,7 +31,8 @@ def int8_matmul(
     weights_dtype: str = None,
 ) -> torch.FloatTensor:
     if quantized_weight_shape is not None:
-        weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8)
+        weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8).t_()
+        scale = scale.t()
     return_dtype = input.dtype
     output_shape = (*input.shape[:-1], weight.shape[-1])
     if svd_up is not None: