From ae4b116f9513ad07b16fcb8649eb4e6f1663ab0c Mon Sep 17 00:00:00 2001 From: Disty0 Date: Mon, 6 Jul 2026 21:40:29 +0300 Subject: [PATCH] dedupe and cleanup sdnq code --- modules/sdnq/dequantizer.py | 198 ++++++++++++++---------------------- modules/sdnq/quant_utils.py | 95 +++++++---------- 2 files changed, 113 insertions(+), 180 deletions(-) diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index 141794cab..fc8967fdf 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -20,8 +20,8 @@ def dequantize_asymmetric( svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None, - dtype: torch.dtype = None, - result_shape: torch.Size = None, + dtype: torch.dtype | None = None, + result_shape: torch.Size | None = None, skip_quantized_matmul: bool = False, re_quantize_for_matmul: bool = False, ) -> torch.FloatTensor: @@ -56,8 +56,8 @@ def dequantize_symmetric( svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None, - dtype: torch.dtype = None, - result_shape: torch.Size = None, + dtype: torch.dtype | None = None, + result_shape: torch.Size | None = None, skip_quantized_matmul: bool = False, re_quantize_for_matmul: bool = False, ) -> torch.FloatTensor: @@ -85,24 +85,29 @@ def dequantize_symmetric( return result -@devices.inference_context() -def dequantize_packed_int_asymmetric(weight: torch.Tensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, weights_dtype: str, svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None, dtype: torch.dtype = None, result_shape: torch.Size = None, skip_quantized_matmul: bool = False, re_quantize_for_matmul: bool = False) -> torch.FloatTensor: - return dequantize_asymmetric(unpack_int(weight, weights_dtype, shape), scale, zero_point, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - - -@devices.inference_context() -def dequantize_packed_int_symmetric(weight: torch.Tensor, scale: torch.FloatTensor, shape: torch.Size, weights_dtype: str, svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None, dtype: torch.dtype = None, result_shape: torch.Size = None, skip_quantized_matmul: bool = False, re_quantize_for_matmul: bool = False) -> torch.FloatTensor: - return dequantize_symmetric(unpack_int(weight, weights_dtype, shape, dtype=scale.dtype), scale, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - - -@devices.inference_context() -def dequantize_packed_float_asymmetric(weight: torch.Tensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, weights_dtype: str, svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None, dtype: torch.dtype = None, result_shape: torch.Size = None, skip_quantized_matmul: bool = False, re_quantize_for_matmul: bool = False) -> torch.FloatTensor: - return dequantize_asymmetric(unpack_float(weight, weights_dtype, shape), scale, zero_point, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - - -@devices.inference_context() -def dequantize_packed_float_symmetric(weight: torch.Tensor, scale: torch.FloatTensor, shape: torch.Size, weights_dtype: str, svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None, dtype: torch.dtype = None, result_shape: torch.Size = None, skip_quantized_matmul: bool = False, re_quantize_for_matmul: bool = False) -> torch.FloatTensor: - return dequantize_symmetric(unpack_float(weight, weights_dtype, shape), scale, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) +def dequantize_weight( + weights_dtype: str, + weight: torch.Tensor, + scale: torch.FloatTensor, + zero_point: torch.FloatTensor | None = None, + svd_up: torch.FloatTensor | None = None, + svd_down: torch.FloatTensor | None = None, + hadamard: torch.FloatTensor | None = None, + dtype: torch.dtype | None = None, + result_shape: torch.Size | None = None, + quantized_weight_shape: torch.Size | None = None, + skip_quantized_matmul: bool = False, + re_quantize_for_matmul: bool = False, +) -> torch.FloatTensor: + if dtype_dict[weights_dtype]["is_packed"]: + if dtype_dict[weights_dtype]["is_integer"]: + weight = unpack_int(weight, weights_dtype, quantized_weight_shape, dtype=scale.dtype) + else: + weight = unpack_float(weight, weights_dtype, quantized_weight_shape) + if dtype_dict[weights_dtype]["is_unsigned"]: + return dequantize_asymmetric(weight, scale, zero_point, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) + else: + return dequantize_symmetric(weight, scale, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) @devices.inference_context() @@ -118,7 +123,7 @@ def re_quantize_int_mm(weight: torch.FloatTensor, matmul_dtype: str = "int8") -> @devices.inference_context() -def re_quantize_uint_mm(weight: torch.FloatTensor, matmul_dtype: str = "uint8") -> tuple[torch.Tensor, torch.FloatTensor]: +def re_quantize_uint_mm(weight: torch.FloatTensor, matmul_dtype: str = "uint8") -> tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]: if weight.ndim > 2: # convs weight = weight.flatten(1,-1) if use_contiguous_mm: @@ -143,9 +148,27 @@ def re_quantize_fp_mm(weight: torch.FloatTensor, matmul_dtype: str = "float8_e4m return weight, scale -@devices.inference_context() -def re_quantize_matmul_asymmetric(weight: torch.Tensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, matmul_dtype: str, result_shape: torch.Size = None, svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None) -> tuple[torch.Tensor, torch.FloatTensor]: - weight = dequantize_asymmetric(weight, scale, zero_point, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=scale.dtype, result_shape=result_shape) +def re_quantize_matmul( + weights_dtype: str, + weight: torch.Tensor, + scale: torch.FloatTensor, + zero_point: torch.FloatTensor | None = None, + svd_up: torch.FloatTensor | None = None, + svd_down: torch.FloatTensor | None = None, + hadamard: torch.FloatTensor | None = None, + matmul_dtype: str = "int8", + result_shape: torch.Size | None = None, + quantized_weight_shape: torch.Size | None = None, +) -> tuple[torch.Tensor, torch.FloatTensor] | tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]: + if dtype_dict[weights_dtype]["is_packed"]: + if dtype_dict[weights_dtype]["is_integer"]: + weight = unpack_int(weight, weights_dtype, quantized_weight_shape, dtype=scale.dtype) + else: + weight = unpack_float(weight, weights_dtype, quantized_weight_shape) + if dtype_dict[weights_dtype]["is_unsigned"]: + weight = dequantize_asymmetric(weight, scale, zero_point, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=scale.dtype, result_shape=result_shape) + else: + weight = dequantize_symmetric(weight, scale, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=scale.dtype, result_shape=result_shape) if dtype_dict[matmul_dtype]["is_integer"]: if dtype_dict[matmul_dtype]["is_unsigned"]: return re_quantize_uint_mm(weight, matmul_dtype=matmul_dtype) @@ -155,38 +178,6 @@ def re_quantize_matmul_asymmetric(weight: torch.Tensor, scale: torch.FloatTensor return re_quantize_fp_mm(weight, matmul_dtype=matmul_dtype) -@devices.inference_context() -def re_quantize_matmul_symmetric(weight: torch.Tensor, scale: torch.FloatTensor, matmul_dtype: str, result_shape: torch.Size = None, svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None) -> tuple[torch.Tensor, torch.FloatTensor]: - weight = dequantize_symmetric(weight, scale, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=scale.dtype, result_shape=result_shape) - if dtype_dict[matmul_dtype]["is_integer"]: - if dtype_dict[matmul_dtype]["is_unsigned"]: - return re_quantize_uint_mm(weight, matmul_dtype=matmul_dtype) - else: - return re_quantize_int_mm(weight, matmul_dtype=matmul_dtype) - else: - return re_quantize_fp_mm(weight, matmul_dtype=matmul_dtype) - - -@devices.inference_context() -def re_quantize_matmul_packed_int_asymmetric(weight: torch.Tensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, weights_dtype: str, matmul_dtype: str, result_shape: torch.Size, svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None) -> tuple[torch.Tensor, torch.FloatTensor]: - return re_quantize_matmul_asymmetric(unpack_int(weight, weights_dtype, shape), scale, zero_point, matmul_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, result_shape=result_shape) - - -@devices.inference_context() -def re_quantize_matmul_packed_int_symmetric(weight: torch.Tensor, scale: torch.FloatTensor, shape: torch.Size, weights_dtype: str, matmul_dtype: str, result_shape: torch.Size = None, svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None) -> tuple[torch.Tensor, torch.FloatTensor]: - return re_quantize_matmul_symmetric(unpack_int(weight, weights_dtype, shape, dtype=scale.dtype), scale, matmul_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, result_shape=result_shape) - - -@devices.inference_context() -def re_quantize_matmul_packed_float_asymmetric(weight: torch.Tensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, weights_dtype: str, matmul_dtype: str, result_shape: torch.Size, svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None) -> tuple[torch.Tensor, torch.FloatTensor]: - return re_quantize_matmul_asymmetric(unpack_float(weight, weights_dtype, shape), scale, zero_point, matmul_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, result_shape=result_shape) - - -@devices.inference_context() -def re_quantize_matmul_packed_float_symmetric(weight: torch.Tensor, scale: torch.FloatTensor, shape: torch.Size, weights_dtype: str, matmul_dtype: str, result_shape: torch.Size = None, svd_up: torch.FloatTensor | None = None, svd_down: torch.FloatTensor | None = None, hadamard: torch.FloatTensor | None = None) -> tuple[torch.Tensor, torch.FloatTensor]: - return re_quantize_matmul_symmetric(unpack_float(weight, weights_dtype, shape), scale, matmul_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, result_shape=result_shape) - - @devices.inference_context() def dequantize_sdnq_module(model: torch.nn.Module) -> torch.nn.Module: if isinstance(model, SDNQLayer): @@ -299,22 +290,18 @@ class SDNQDequantizer: ) -> tuple[torch.Tensor, torch.FloatTensor]: # pylint: disable=unused-argument if hadamard is None and self.use_hadamard and not non_hadamard: hadamard = get_hadamard(self.hadamard_group_size, dtype=self.result_dtype, device=weight.device) - if self.is_packed: - if self.is_integer: - if self.is_unsigned: - return re_quantize_matmul_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.weights_dtype, self.quantized_matmul_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, result_shape=self.result_shape) - else: - return re_quantize_matmul_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.weights_dtype, self.quantized_matmul_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, result_shape=self.result_shape) - else: - if self.is_unsigned: - return re_quantize_matmul_packed_float_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.weights_dtype, self.quantized_matmul_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, result_shape=self.result_shape) - else: - return re_quantize_matmul_packed_float_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.weights_dtype, self.quantized_matmul_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, result_shape=self.result_shape) - else: - if self.is_unsigned: - return re_quantize_matmul_asymmetric_compiled(weight, scale, zero_point, self.quantized_matmul_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, result_shape=self.result_shape) - else: - return re_quantize_matmul_symmetric_compiled(weight, scale, self.quantized_matmul_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, result_shape=self.result_shape) + return re_quantize_matmul_compiled( + self.weights_dtype, + weight, + scale, + zero_point=zero_point, + svd_up=svd_up, + svd_down=svd_down, + hadamard=hadamard, + matmul_dtype=self.quantized_matmul_dtype, + result_shape=self.result_shape, + quantized_weight_shape=self.quantized_weight_shape, + ) @devices.inference_context() def __call__( @@ -328,60 +315,33 @@ class SDNQDequantizer: skip_quantized_matmul: bool = False, non_hadamard: bool = False, skip_compile: bool = False, - dtype: torch.dtype = None, + dtype: torch.dtype | None = None, ) -> torch.FloatTensor: # pylint: disable=unused-argument if dtype is None: dtype = self.result_dtype if hadamard is None and self.use_hadamard and not non_hadamard: hadamard = get_hadamard(self.hadamard_group_size, dtype=dtype, device=weight.device) re_quantize_for_matmul = self.re_quantize_for_matmul or self.is_packed - if self.is_packed: - if self.is_integer: - if self.is_unsigned: - if skip_compile: # compiled training needs to be traced with the original function - return dequantize_packed_int_asymmetric(weight, scale, zero_point, self.quantized_weight_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - return dequantize_packed_int_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - if skip_compile: - return dequantize_packed_int_symmetric(weight, scale, self.quantized_weight_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - return dequantize_packed_int_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - if self.is_unsigned: - if skip_compile: # compiled training needs to be traced with the original function - return dequantize_packed_float_asymmetric(weight, scale, zero_point, self.quantized_weight_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - return dequantize_packed_float_asymmetric_compiled(weight, scale, zero_point, self.quantized_weight_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - if skip_compile: - return dequantize_packed_float_symmetric(weight, scale, self.quantized_weight_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - return dequantize_packed_float_symmetric_compiled(weight, scale, self.quantized_weight_shape, self.weights_dtype, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - if self.is_unsigned: - if skip_compile: - return dequantize_asymmetric(weight, scale, zero_point, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - return dequantize_asymmetric_compiled(weight, scale, zero_point, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - if skip_compile: - return dequantize_symmetric(weight, scale, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) - else: - return dequantize_symmetric_compiled(weight, scale, svd_up=svd_up, svd_down=svd_down, hadamard=hadamard, dtype=dtype, result_shape=self.result_shape, skip_quantized_matmul=skip_quantized_matmul, re_quantize_for_matmul=re_quantize_for_matmul) + dequantize_weight_func = dequantize_weight if skip_compile else dequantize_weight_compiled + return dequantize_weight_func( + self.weights_dtype, + weight, + scale, + zero_point=zero_point, + svd_up=svd_up, + svd_down=svd_down, + hadamard=hadamard, + dtype=dtype, + result_shape=self.result_shape, + quantized_weight_shape=self.quantized_weight_shape, + skip_quantized_matmul=skip_quantized_matmul, + re_quantize_for_matmul=re_quantize_for_matmul, + ) dequantize_asymmetric_compiled = compile_func(dequantize_asymmetric) dequantize_symmetric_compiled = compile_func(dequantize_symmetric) -dequantize_packed_int_asymmetric_compiled = compile_func(dequantize_packed_int_asymmetric) -dequantize_packed_int_symmetric_compiled = compile_func(dequantize_packed_int_symmetric) -dequantize_packed_float_asymmetric_compiled = compile_func(dequantize_packed_float_asymmetric) -dequantize_packed_float_symmetric_compiled = compile_func(dequantize_packed_float_symmetric) -re_quantize_matmul_asymmetric_compiled = compile_func(re_quantize_matmul_asymmetric) -re_quantize_matmul_symmetric_compiled = compile_func(re_quantize_matmul_symmetric) -re_quantize_matmul_packed_int_asymmetric_compiled = compile_func(re_quantize_matmul_packed_int_asymmetric) -re_quantize_matmul_packed_int_symmetric_compiled = compile_func(re_quantize_matmul_packed_int_symmetric) -re_quantize_matmul_packed_float_asymmetric_compiled = compile_func(re_quantize_matmul_packed_float_asymmetric) -re_quantize_matmul_packed_float_symmetric_compiled = compile_func(re_quantize_matmul_packed_float_symmetric) +dequantize_weight_compiled = compile_func(dequantize_weight) +re_quantize_matmul_compiled = compile_func(re_quantize_matmul) torch.serialization.add_safe_globals([SDNQDequantizer]) diff --git a/modules/sdnq/quant_utils.py b/modules/sdnq/quant_utils.py index 5ea676fac..cac1e34ce 100644 --- a/modules/sdnq/quant_utils.py +++ b/modules/sdnq/quant_utils.py @@ -8,32 +8,32 @@ from .common import dtype_dict, use_contiguous_mm, conv_types, conv_transpose_ty @devices.inference_context() -def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: int | list[int], weights_dtype: str) -> tuple[torch.FloatTensor, torch.FloatTensor]: - zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True) - scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"]) +def get_scale_asymmetric(weight: torch.FloatTensor, dim: int | list[int], weights_dtype: str) -> tuple[torch.FloatTensor, torch.FloatTensor]: + zero_point, scale = torch.aminmax(weight, dim=dim, keepdims=True) + scale = scale.sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"]) if dtype_dict[weights_dtype]["min"] != 0: - zero_point.sub_(torch.mul(scale, dtype_dict[weights_dtype]["min"])) + zero_point.sub_(scale, alpha=dtype_dict[weights_dtype]["min"]) return scale, zero_point @devices.inference_context() -def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: int | list[int], weights_dtype: str) -> torch.FloatTensor: - return torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).div_(dtype_dict[weights_dtype]["max"]) +def get_scale_symmetric(weight: torch.FloatTensor, dim: int | list[int], weights_dtype: str) -> torch.FloatTensor: + return torch.amax(weight.abs(), dim=dim, keepdims=True).div_(dtype_dict[weights_dtype]["max"]) @devices.inference_context() -def quantize_weight(weight: torch.FloatTensor, reduction_axes: int | list[int], weights_dtype: str, dtype: torch.dtype = None, use_stochastic_rounding: bool = False) -> tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]: +def quantize_weight(weight: torch.FloatTensor, dim: int | list[int], weights_dtype: str, dtype: torch.dtype = None, use_stochastic_rounding: bool = False) -> tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]: if weight.dtype != torch.float64: weight = weight.to(dtype=torch.float32, copy=False) if dtype_dict[weights_dtype]["is_unsigned"]: - scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype) + scale, zero_point = get_scale_asymmetric(weight, dim, weights_dtype) if dtype is not None: scale = scale.to(dtype=dtype) zero_point = zero_point.to(dtype=dtype) quantized_weight = torch.sub(weight, zero_point).div_(scale) else: - scale = get_scale_symmetric(weight, reduction_axes, weights_dtype) + scale = get_scale_symmetric(weight, dim, weights_dtype) zero_point = None if dtype is not None: scale = scale.to(dtype=dtype) @@ -188,65 +188,38 @@ def prepare_svd_for_matmul(svd_up: torch.FloatTensor, svd_down: torch.FloatTenso @devices.inference_context() -def quantize_int_mm(input: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "int8") -> tuple[torch.Tensor, torch.FloatTensor]: +def quantize_int_mm(weight: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "int8", use_sr: bool = False) -> tuple[torch.Tensor, torch.FloatTensor]: if hadamard is not None: - input = rotate_hadamard(input, hadamard=hadamard) - scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(dtype_dict[matmul_dtype]["max"]) - input = torch.div(input, scale).round_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"]) - return input, scale + weight = rotate_hadamard(weight, hadamard=hadamard) + scale = get_scale_symmetric(weight, dim, matmul_dtype) + weight = torch.div(weight, scale) + if use_sr: + weight = weight.add_(torch.randn_like(weight), alpha=0.1) + weight = weight.round_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"]) + return weight, scale @devices.inference_context() -def quantize_int_mm_sr(input: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "int8") -> tuple[torch.Tensor, torch.FloatTensor]: +def quantize_uint_mm(weight: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "uint8", use_sr: bool = False) -> tuple[torch.FloatTensor, torch.FloatTensor]: if hadamard is not None: - input = rotate_hadamard(input, hadamard=hadamard) - scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(dtype_dict[matmul_dtype]["max"]) - 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 - - -@devices.inference_context() -def quantize_uint_mm(input: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "uint8") -> tuple[torch.FloatTensor, torch.FloatTensor]: - if hadamard is not None: - input = rotate_hadamard(input, hadamard=hadamard) + weight = rotate_hadamard(weight, hadamard=hadamard) matmul_dtype = matmul_dtype.removeprefix("u") - zero_point = torch.amin(input, dim=dim, keepdims=True) - scale = torch.amax(input, dim=dim, keepdims=True).sub_(zero_point).div_(dtype_dict[matmul_dtype]["max"] - dtype_dict[matmul_dtype]["min"]) - if dtype_dict[matmul_dtype]["min"] != 0: - zero_point.sub_(scale, alpha=dtype_dict[matmul_dtype]["min"]) - input = torch.sub(input, zero_point).div_(scale).round_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"]) - return input, scale, zero_point + scale, zero_point = get_scale_asymmetric(weight, dim, matmul_dtype) + weight = torch.sub(weight, zero_point).div_(scale) + if use_sr: + weight = weight.add_(torch.randn_like(weight), alpha=0.1) + weight = weight.round_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"]) + return weight, scale, zero_point @devices.inference_context() -def quantize_uint_mm_sr(input: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "uint8") -> tuple[torch.FloatTensor, torch.FloatTensor]: +def quantize_fp_mm(weight: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "float8_e4m3fn", use_sr: bool = False) -> tuple[torch.Tensor, torch.FloatTensor]: if hadamard is not None: - input = rotate_hadamard(input, hadamard=hadamard) - matmul_dtype = matmul_dtype.removeprefix("u") - zero_point = torch.amin(input, dim=dim, keepdims=True) - scale = torch.amax(input, dim=dim, keepdims=True).sub_(zero_point).div_(dtype_dict[matmul_dtype]["max"] - dtype_dict[matmul_dtype]["min"]) - if dtype_dict[matmul_dtype]["min"] != 0: - zero_point.sub_(scale, alpha=dtype_dict[matmul_dtype]["min"]) - input = torch.sub(input, zero_point).div_(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, zero_point - - -@devices.inference_context() -def quantize_fp_mm(input: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "float8_e4m3fn") -> tuple[torch.Tensor, torch.FloatTensor]: - if hadamard is not None: - input = rotate_hadamard(input, hadamard=hadamard) - scale = torch.amax(input.abs(), dim=dim, keepdims=True).div_(dtype_dict[matmul_dtype]["max"]) - input = torch.div(input, scale).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 - - -@devices.inference_context() -def quantize_fp_mm_sr(input: torch.FloatTensor, dim: int = -1, hadamard: torch.FloatTensor | None = None, matmul_dtype: str = "float8_e4m3fn") -> tuple[torch.Tensor, torch.FloatTensor]: - if hadamard is not None: - input = rotate_hadamard(input, hadamard=hadamard) - 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)).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 + weight = rotate_hadamard(weight, hadamard=hadamard) + scale = get_scale_symmetric(weight, dim, matmul_dtype) + if use_sr: + mantissa_difference = 1 << (23 - dtype_dict[matmul_dtype]["mantissa"]) + weight = weight.to(dtype=torch.float32).view(dtype=torch.int32) + weight = weight.add_(torch.randint_like(weight, low=0, high=mantissa_difference, dtype=torch.int32)).bitwise_and_(-mantissa_difference).view(dtype=torch.float32) + weight = torch.div(weight, scale).nan_to_num_().clamp_(dtype_dict[matmul_dtype]["min"], dtype_dict[matmul_dtype]["max"]).to(dtype=dtype_dict[matmul_dtype]["torch_dtype"]) + return weight, scale