mirror of
https://github.com/vladmandic/automatic
synced 2026-09-02 19:10:46 +02:00
dedupe and cleanup sdnq code
This commit is contained in:
Disty0
+79
-119
@@ -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])
|
||||
|
||||
+34
-61
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user