diff --git a/modules/sdnq/layers/conv/conv_fp16.py b/modules/sdnq/layers/conv/conv_fp16.py index 71ddf3162..d81750895 100644 --- a/modules/sdnq/layers/conv/conv_fp16.py +++ b/modules/sdnq/layers/conv/conv_fp16.py @@ -5,6 +5,7 @@ from typing import List import torch from ...common import compile_func, fp_mm_func # noqa: TID252 +from ...packed_float import unpack_float # noqa: TID252 from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 from .forward import get_conv_args, process_conv_input @@ -24,6 +25,8 @@ def conv_fp16_matmul( bias: torch.FloatTensor = None, svd_up: torch.FloatTensor = None, svd_down: torch.FloatTensor = None, + quantized_weight_shape: torch.Size = None, + weights_dtype: str = None, ) -> torch.FloatTensor: return_dtype = input.dtype input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) @@ -35,7 +38,10 @@ def conv_fp16_matmul( 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") - weight = weight.to(dtype=torch.float16) # fp8 weights + if quantized_weight_shape is not None: + weight = unpack_float(weight, quantized_weight_shape, weights_dtype).to(dtype=torch.float16) + elif weight.dtype != torch.float16: + weight = weight.to(dtype=torch.float16) # fp8 weights input, weight = check_mats(input, weight) if groups == 1: @@ -64,8 +70,10 @@ def conv_fp16_matmul( def quantized_conv_forward_fp16_matmul(self, input) -> torch.FloatTensor: if self.sdnq_dequantizer.re_quantize_for_matmul: weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None) + quantized_weight_shape = None else: weight, scale = self.weight, self.scale + quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) return conv_fp16_matmul( input, weight, scale, @@ -76,6 +84,8 @@ def quantized_conv_forward_fp16_matmul(self, input) -> torch.FloatTensor: bias=self.bias, svd_up=self.svd_up, svd_down=self.svd_down, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=self.sdnq_dequantizer.weights_dtype, ) diff --git a/modules/sdnq/layers/conv/conv_fp8.py b/modules/sdnq/layers/conv/conv_fp8.py index 4fcad6509..738fb54f7 100644 --- a/modules/sdnq/layers/conv/conv_fp8.py +++ b/modules/sdnq/layers/conv/conv_fp8.py @@ -5,6 +5,7 @@ from typing import List import torch from ...common import compile_func # noqa: TID252 +from ...packed_float import unpack_float # noqa: TID252 from .forward import get_conv_args, process_conv_input from ..linear.linear_fp8 import quantize_fp_mm_input # noqa: TID252 @@ -23,6 +24,8 @@ def conv_fp8_matmul( bias: torch.FloatTensor = None, svd_up: torch.FloatTensor = None, svd_down: torch.FloatTensor = None, + quantized_weight_shape: torch.Size = None, + weights_dtype: str = None, ) -> torch.FloatTensor: return_dtype = input.dtype input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) @@ -31,6 +34,8 @@ def conv_fp8_matmul( 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) input, weight = check_mats(input, weight) if groups == 1: @@ -71,8 +76,10 @@ def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor: return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias) if self.sdnq_dequantizer.re_quantize_for_matmul: weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None) + quantized_weight_shape = None else: weight, scale = self.weight, self.scale + quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) return conv_fp8_matmul( input, weight, scale, @@ -83,6 +90,8 @@ def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor: bias=self.bias, svd_up=self.svd_up, svd_down=self.svd_down, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=self.sdnq_dequantizer.weights_dtype, ) diff --git a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py index 2079bea33..c8ad77e3e 100644 --- a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py +++ b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py @@ -5,6 +5,7 @@ from typing import List import torch from ...common import compile_func # noqa: TID252 +from ...packed_float import unpack_float # noqa: TID252 from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 from .forward import get_conv_args, process_conv_input @@ -24,6 +25,8 @@ def conv_fp8_matmul_tensorwise( bias: torch.FloatTensor = None, svd_up: torch.FloatTensor = None, svd_down: torch.FloatTensor = None, + quantized_weight_shape: torch.Size = None, + weights_dtype: str = None, ) -> torch.FloatTensor: return_dtype = input.dtype input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) @@ -35,6 +38,8 @@ def conv_fp8_matmul_tensorwise( 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) input, weight = check_mats(input, weight) dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) @@ -66,8 +71,10 @@ def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTens return self._conv_forward(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias) if self.sdnq_dequantizer.re_quantize_for_matmul: weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None) + quantized_weight_shape = None else: weight, scale = self.weight, self.scale + quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) return conv_fp8_matmul_tensorwise( input, weight, scale, @@ -78,6 +85,8 @@ def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTens bias=self.bias, svd_up=self.svd_up, svd_down=self.svd_down, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=self.sdnq_dequantizer.weights_dtype, ) diff --git a/modules/sdnq/layers/conv/conv_int8.py b/modules/sdnq/layers/conv/conv_int8.py index 9eaee44c7..7091f987e 100644 --- a/modules/sdnq/layers/conv/conv_int8.py +++ b/modules/sdnq/layers/conv/conv_int8.py @@ -15,18 +15,18 @@ from ..linear.forward import check_mats # noqa: TID252 def conv_int8_matmul( input: torch.FloatTensor, - weight: torch.CharTensor, - bias: torch.FloatTensor, + weight: torch.Tensor, scale: torch.FloatTensor, - svd_up: torch.FloatTensor, - svd_down: torch.FloatTensor, - quantized_weight_shape: torch.Size, result_shape: torch.Size, - weights_dtype: str, reversed_padding_repeated_twice: List[int], padding_mode: str, conv_type: int, groups: int, stride: List[int], padding: List[int], dilation: List[int], + bias: torch.FloatTensor = None, + svd_up: torch.FloatTensor = None, + svd_down: torch.FloatTensor = None, + quantized_weight_shape: torch.Size = None, + weights_dtype: str = None, ) -> torch.FloatTensor: return_dtype = input.dtype input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) @@ -73,18 +73,19 @@ def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor: weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None) quantized_weight_shape = None else: - weight = self.weight - scale = self.scale + weight, scale = self.weight, self.scale quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None return conv_int8_matmul( - input, weight, self.bias, - scale, self.svd_up, self.svd_down, - quantized_weight_shape, + input, weight, scale, self.sdnq_dequantizer.result_shape, - self.sdnq_dequantizer.weights_dtype, self._reversed_padding_repeated_twice, self.padding_mode, conv_type, self.groups, stride, padding, dilation, + bias=self.bias, + svd_up=self.svd_up, + svd_down=self.svd_down, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=self.sdnq_dequantizer.weights_dtype, ) diff --git a/modules/sdnq/layers/linear/linear_fp16.py b/modules/sdnq/layers/linear/linear_fp16.py index db8a5cab0..d8dc7b9c1 100644 --- a/modules/sdnq/layers/linear/linear_fp16.py +++ b/modules/sdnq/layers/linear/linear_fp16.py @@ -3,6 +3,7 @@ import torch from ...common import compile_func, fp_mm_func # noqa: TID252 +from ...packed_float import unpack_float # noqa: TID252 from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 from .forward import check_mats @@ -16,7 +17,13 @@ def fp16_matmul( bias: torch.FloatTensor = None, svd_up: torch.FloatTensor = None, svd_down: torch.FloatTensor = None, + quantized_weight_shape: torch.Size = None, + 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) + elif weight.dtype != torch.float16: + weight = weight.to(dtype=torch.float16) # fp8 weights return_dtype = input.dtype output_shape = (*input.shape[:-1], weight.shape[-1]) if svd_up is not None: @@ -26,7 +33,6 @@ def 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") - weight = weight.to(dtype=torch.float16) # fp8 weights input, weight = check_mats(input, weight) if bias is not None: return dequantize_symmetric_with_bias(fp_mm_func(input, weight), scale, bias, dtype=return_dtype, result_shape=output_shape) @@ -37,9 +43,18 @@ def fp16_matmul( def quantized_linear_forward_fp16_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: if self.sdnq_dequantizer.re_quantize_for_matmul: weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None) + quantized_weight_shape = None else: weight, scale = self.weight, self.scale - return fp16_matmul(input, weight, scale, bias=self.bias, svd_up=self.svd_up, svd_down=self.svd_down) + quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None + return fp16_matmul( + input, weight, scale, + bias=self.bias, + svd_up=self.svd_up, + svd_down=self.svd_down, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=self.sdnq_dequantizer.weights_dtype, + ) fp16_matmul = compile_func(fp16_matmul) diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py index d8f65ad6f..db1f29c9c 100644 --- a/modules/sdnq/layers/linear/linear_fp8.py +++ b/modules/sdnq/layers/linear/linear_fp8.py @@ -5,6 +5,7 @@ from typing import Tuple import torch from ...common import compile_func # noqa: TID252 +from ...packed_float import unpack_float # noqa: TID252 from ...dequantizer import quantize_fp_mm # noqa: TID252 from .forward import check_mats @@ -23,7 +24,11 @@ def fp8_matmul( bias: torch.FloatTensor = None, svd_up: torch.FloatTensor = None, svd_down: torch.FloatTensor = None, + quantized_weight_shape: torch.Size = None, + 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) return_dtype = input.dtype output_shape = (*input.shape[:-1], weight.shape[-1]) if svd_up is not None: @@ -45,9 +50,18 @@ def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias) if self.sdnq_dequantizer.re_quantize_for_matmul: weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None) + quantized_weight_shape = None else: weight, scale = self.weight, self.scale - return fp8_matmul(input, weight, scale, bias=self.bias, svd_up=self.svd_up, svd_down=self.svd_down) + quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None + return fp8_matmul( + input, weight, scale, + bias=self.bias, + svd_up=self.svd_up, + svd_down=self.svd_down, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=self.sdnq_dequantizer.weights_dtype, + ) fp8_matmul = compile_func(fp8_matmul) diff --git a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py index 14560db48..53f53da07 100644 --- a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py +++ b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py @@ -5,6 +5,7 @@ from typing import Tuple import torch from ...common import compile_func # noqa: TID252 +from ...packed_float import unpack_float # noqa: TID252 from ...dequantizer import quantize_fp_mm, dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 from .forward import check_mats @@ -26,7 +27,11 @@ def fp8_matmul_tensorwise( bias: torch.FloatTensor = None, svd_up: torch.FloatTensor = None, svd_down: torch.FloatTensor = None, + quantized_weight_shape: torch.Size = None, + 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) return_dtype = input.dtype output_shape = (*input.shape[:-1], weight.shape[-1]) if svd_up is not None: @@ -49,9 +54,18 @@ def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTenso return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, self.zero_point, self.svd_up, self.svd_down, skip_quantized_matmul=True), self.bias) if self.sdnq_dequantizer.re_quantize_for_matmul: weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None) + quantized_weight_shape = None else: weight, scale = self.weight, self.scale - return fp8_matmul_tensorwise(input, weight, scale, bias=self.bias, svd_up=self.svd_up, svd_down=self.svd_down) + quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None + return fp8_matmul_tensorwise( + input, weight, scale, + bias=self.bias, + svd_up=self.svd_up, + svd_down=self.svd_down, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=self.sdnq_dequantizer.weights_dtype, + ) fp8_matmul_tensorwise = compile_func(fp8_matmul_tensorwise) diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py index 14efcea34..05b06e0e2 100644 --- a/modules/sdnq/layers/linear/linear_int8.py +++ b/modules/sdnq/layers/linear/linear_int8.py @@ -55,8 +55,7 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, self.zero_point, None, None) quantized_weight_shape = None else: - weight = self.weight - scale = self.scale + weight, scale = self.weight, self.scale quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None return int8_matmul( input, weight, scale, @@ -64,7 +63,7 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc svd_up=self.svd_up, svd_down=self.svd_down, quantized_weight_shape=quantized_weight_shape, - weights_dtype=self.sdnq_dequantizer.weights_dtype + weights_dtype=self.sdnq_dequantizer.weights_dtype, ) diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index 6a6a0eb6e..f0a7d2426 100644 --- a/modules/sdnq/quantizer.py +++ b/modules/sdnq/quantizer.py @@ -228,7 +228,15 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int dtype_dict[weights_dtype]["is_unsigned"] or dtype_dict[weights_dtype]["is_integer"] != dtype_dict[quantized_matmul_dtype]["is_integer"] or dtype_dict[weights_dtype]["num_bits"] > dtype_dict[quantized_matmul_dtype]["num_bits"] - or (dtype_dict[weights_dtype]["is_packed"] and not dtype_dict[weights_dtype]["is_integer"]) + or ( + dtype_dict[weights_dtype]["is_packed"] + and not dtype_dict[weights_dtype]["is_integer"] + and not dtype_dict[quantized_matmul_dtype]["is_integer"] + and ( + dtype_dict[weights_dtype]["num_bits"] >= dtype_dict[quantized_matmul_dtype]["num_bits"] + or dtype_dict[weights_dtype]["max"] > dtype_dict[quantized_matmul_dtype]["max"] + ) + ) ) if layer_class_name in conv_types: @@ -348,7 +356,7 @@ def sdnq_quantize_layer_weight(weight, layer_class_name=None, weights_dtype="int scale.t_() weight.t_() weight = prepare_weight_for_matmul(weight) - if not use_tensorwise_fp8_matmul and not dtype_dict[weights_dtype]["is_integer"]: + if not use_tensorwise_fp8_matmul and not dtype_dict[quantized_matmul_dtype]["is_integer"]: scale = scale.to(dtype=torch.float32) sdnq_dequantizer = SDNQDequantizer( @@ -956,8 +964,9 @@ class SDNQConfig(QuantizationConfigMixin): Args: weights_dtype (`str`, *optional*, defaults to `"int8"`): - The target dtype for the weights after quantization. Supported values are: - ("int16", "int8", "int7", "int6", "int5", "int4", "int3", "int2", "uint16", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float16", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz") + The target dtype for the weights after quantization. + Check out `sdnq.common.accepted_weight_dtypes` for all the supported values. + These are some of the recommended values to use: ("int8", "int7", "int6", "uint5", "uint4", "uint3", "uint2", "float8_e4m3fn", "float7_e3m3fn", "float6_e3m2fn", "float5_e2m2fn", "float4_e2m1fn", "float3_e1m1fn", "float2_e1m0fn") quantized_matmul_dtype (`str`, *optional*, defaults to `None`): The target dtype for quantized matmul. `None` will use "int8" with integer weight dtypes and "float8_e4m3fn" or "float16" with float weight dtypes.