SDNQ add SVD support for Convs

modules/sdnq/dequantizer.py

@@ -15,7 +15,10 @@ def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, ze
     if svd_up is not None:
         if skip_quantized_matmul:
             svd_up, svd_down = svd_up.t(), svd_down.t()
-        result = torch.addmm(result, svd_up, svd_down)
+        if result.ndim > 2 and weight.ndim > 2: # convs
+            result = result.add_(torch.mm(svd_up, svd_down).unflatten(-1, (*result.shape[1:],)))
+        else:
+            result = result.addmm_(svd_up, svd_down)
     if dtype is not None:
         result = result.to(dtype=dtype)
     return result
@@ -30,7 +33,10 @@ def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dty
     if svd_up is not None:
         if skip_quantized_matmul:
             svd_up, svd_down = svd_up.t(), svd_down.t()
-        result = torch.addmm(result, svd_up, svd_down)
+        if result.ndim > 2 and weight.ndim > 2: # convs
+            result = result.add_(torch.mm(svd_up, svd_down).unflatten(-1, (*result.shape[1:],)))
+        else:
+            result = result.addmm_(svd_up, svd_down)
     if dtype is not None:
         result = result.to(dtype=dtype)
     return result

modules/sdnq/layers/conv/conv_fp8.py

@@ -34,7 +34,7 @@ def conv_fp8_matmul(
     if groups == 1:
         if bias is not None and bias.dtype != torch.bfloat16:
             bias = bias.to(dtype=torch.bfloat16)
-        result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16).view(mm_output_shape).to(return_dtype)
+        result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16)
     else:
         scale = scale.view(groups, 1, scale.shape[1] // groups)
         input_scale = input_scale.view(groups, input_scale.shape[0] // groups, 1)
@@ -50,9 +50,10 @@ def conv_fp8_matmul(
         else:
             for i in range(groups):
                 result.append(torch._scaled_mm(input[:, i], weight[:, i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=torch.bfloat16))
-        result = torch.cat(result, dim=-1).view(mm_output_shape).to(return_dtype)
+        result = torch.cat(result, dim=-1)
     if svd_up is not None:
         result.add_(svd_bias)
+    result = result.view(mm_output_shape).to(return_dtype)
 
     if conv_type == 1:
         result = result.transpose_(1,2)

modules/sdnq/layers/linear/linear_fp8.py

@@ -31,9 +31,10 @@ def fp8_matmul(
     input, weight = check_mats(input, weight)
     if bias is not None and bias.dtype != torch.bfloat16:
         bias = bias.to(dtype=torch.bfloat16)
-    result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16).view(output_shape).to(return_dtype)
+    result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16)
     if svd_up is not None:
         result.add_(svd_bias)
+    result = result.view(output_shape).to(return_dtype)
     return result
 
 

modules/sdnq/quantizer.py

@@ -49,10 +49,18 @@ def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[i
 
 
 def apply_svdquant(weight: torch.FloatTensor, rank: int = 32) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
+    reshape_weight = False
+    if weight.ndim > 2: # convs
+        reshape_weight = True
+        weight_shape = weight.shape
+        weight = weight.flatten(1,-1)
     U, S, svd_down = torch.svd_lowrank(weight, q=rank)
     svd_up = torch.mul(U, S.unsqueeze(0))
     svd_down = svd_down.t_()
-    return weight.sub_(torch.mm(svd_up, svd_down)), svd_up, svd_down
+    weight = weight.sub_(torch.mm(svd_up, svd_down))
+    if reshape_weight:
+        weight = weight.unflatten(-1, (*weight_shape[1:],))
+    return weight, svd_up, svd_down
 
 
 @devices.inference_context()
@@ -118,13 +126,12 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
         if layer.weight.dtype != torch.float32:
             layer.weight.data = layer.weight.to(dtype=torch.float32)
 
-        if use_svd and is_linear_type:
+        if use_svd:
             layer.weight.data, svd_up, svd_down = apply_svdquant(layer.weight, rank=svd_rank)
             if use_quantized_matmul:
                 svd_up = svd_up.t_()
                 svd_down = svd_down.t_()
         else:
-            use_svd = False
             svd_up, svd_down = None, None
 
         if group_size == 0: