From 3c8be0f55f2c82771d976ea5f7a74f387261a797 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Fri, 30 May 2025 04:47:29 +0300 Subject: [PATCH] SDNQ add uint2 --- CHANGELOG.md | 2 +- modules/model_quant_sdnq.py | 72 +++++++++++++++++++++++++++++++++++-- modules/shared.py | 2 +- 3 files changed, 72 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index aa601b1dc..95be9e83a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ - `INT8` -> `uint8` - `INT4_SYM` -> `int4` - `INT4` -> `uint4` - - Add `float8_e4m3fn` and `float8_e5m2` support + - Add `float8_e4m3fn`, `float8_e5m2` and `uint2` support - Add quantized matmul support for `float8_e4m3fn` - Set the default quant mode to `pre` - Use per token input quant with int8 and fp8 quantized matmul diff --git a/modules/model_quant_sdnq.py b/modules/model_quant_sdnq.py index be1aa7efe..87da6b592 100644 --- a/modules/model_quant_sdnq.py +++ b/modules/model_quant_sdnq.py @@ -20,6 +20,7 @@ dtype_dict = { "uint8": {"min": 0, "max": 255, "num_bits": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, "int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, "uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, + "uint2": {"min": 0, "max": 3, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, "float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False}, "float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False}, "float8_e4m3fnuz": {"min": -240, "max": 240, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False}, @@ -71,9 +72,12 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]: use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0 - if not use_quantized_matmul and (group_size > 0 or (dtype_dict[weights_dtype]["num_bits"] == 4 and group_size != -1)): + if not use_quantized_matmul and (group_size > 0 or (dtype_dict[weights_dtype]["num_bits"] < 8 and group_size != -1)): if group_size == 0: - group_size = 64 + if dtype_dict[weights_dtype]["num_bits"] < 4: + group_size = 32 + else: + group_size = 64 num_of_groups = channel_size // group_size if group_size >= channel_size: @@ -258,6 +262,10 @@ def decompress_int4_symmetric(input: torch.Tensor, scale: torch.Tensor, shape: t return decompress_symmetric(unpack_int4(input, shape, dtype=scale.dtype), scale, dtype, result_shape) +def decompress_int2_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size) -> torch.Tensor: + return decompress_asymmetric(unpack_uint2(input, shape), scale, zero_point, dtype, result_shape) + + def pack_uint4(tensor: torch.Tensor) -> torch.Tensor: if tensor.dtype != torch.uint8: raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.") @@ -272,6 +280,23 @@ def pack_int4(tensor: torch.Tensor) -> torch.Tensor: return pack_uint4((tensor + 8).to(dtype=torch.uint8)) +def pack_uint2(tensor: torch.Tensor) -> torch.Tensor: + if tensor.dtype != torch.uint8: + raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.") + packed_tensor = tensor.contiguous().reshape(-1, 4) + packed_tensor = torch.bitwise_or( + torch.bitwise_or( + torch.bitwise_and(packed_tensor[:, 0], 3), + torch.bitwise_left_shift(torch.bitwise_and(packed_tensor[:, 1], 3), 2) + ), + torch.bitwise_or( + torch.bitwise_left_shift(torch.bitwise_and(packed_tensor[:, 2], 3), 4), + torch.bitwise_left_shift(packed_tensor[:, 3], 6) + ), + ) + return packed_tensor + + def unpack_uint4(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).reshape(shape) return result @@ -284,6 +309,19 @@ def unpack_int4(packed_tensor: torch.Tensor, shape: torch.Size, dtype: Optional[ return result +def unpack_uint2(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: + result = torch.stack( + ( + torch.bitwise_and(packed_tensor, 3), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 2), 3), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 4), 3), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 6), 3), + ), + dim=-1 + ).reshape(shape) + return result + + def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.FloatTensor, torch.FloatTensor]: input = input.flatten(0,-2).contiguous() input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 448) @@ -499,11 +537,38 @@ class INT4SymmetricWeightsDecompressor(torch.nn.Module): return decompress_int4_symmetric_compiled(weight, self.scale, self.compressed_weight_shape, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul) +class INT2AsymmetricWeightsDecompressor(torch.nn.Module): + def __init__( + self, + scale: torch.Tensor, + zero_point: torch.Tensor, + compressed_weight_shape: torch.Size, + result_dtype: torch.dtype, + result_shape: torch.Size, + **kwargs, + ): + super().__init__() + self.weights_dtype = "uint2" + self.use_quantized_matmul = False + self.compressed_weight_shape = compressed_weight_shape + self.result_dtype = result_dtype + self.result_shape = result_shape + self.register_buffer("scale", scale) + self.register_buffer("zero_point", zero_point) + + def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: + return pack_uint2(weight.to(dtype=torch.uint8)) + + def forward(self, weight, **kwargs): + return decompress_int2_asymmetric_compiled(weight, self.scale, self.zero_point, self.compressed_weight_shape, self.result_dtype, self.result_shape) + + decompressor_dict = { "int8": SymmetricWeightsDecompressor, "uint8": AsymmetricWeightsDecompressor, "int4": INT4SymmetricWeightsDecompressor, "uint4": INT4AsymmetricWeightsDecompressor, + "uint2": INT2AsymmetricWeightsDecompressor, "float8_e4m3fn": SymmetricWeightsDecompressor, "float8_e4m3fnuz": SymmetricWeightsDecompressor, "float8_e5m2": SymmetricWeightsDecompressor, @@ -716,6 +781,7 @@ if shared.opts.sdnq_decompress_compile: decompress_symmetric_compiled = torch.compile(decompress_symmetric, fullgraph=True) decompress_int4_asymmetric_compiled = torch.compile(decompress_int4_asymmetric, fullgraph=True) decompress_int4_symmetric_compiled = torch.compile(decompress_int4_symmetric, fullgraph=True) + decompress_int2_asymmetric_compiled = torch.compile(decompress_int2_asymmetric, fullgraph=True) fp8_matmul = torch.compile(fp8_matmul, fullgraph=True) fp8_matmul_sm89 = torch.compile(fp8_matmul_sm89, fullgraph=True) if devices.backend != "ipex": # pytorch uses the cpu device in torch._int_mm op with ipex + torch.compile @@ -731,6 +797,7 @@ if shared.opts.sdnq_decompress_compile: decompress_symmetric_compiled = decompress_symmetric decompress_int4_asymmetric_compiled = decompress_int4_asymmetric decompress_int4_symmetric_compiled = decompress_int4_symmetric + decompress_int2_asymmetric_compiled = decompress_int2_asymmetric quantize_int8_matmul_input_compiled = quantize_int8_matmul_input unpack_int4_compiled = unpack_int4 else: @@ -738,5 +805,6 @@ else: decompress_symmetric_compiled = decompress_symmetric decompress_int4_asymmetric_compiled = decompress_int4_asymmetric decompress_int4_symmetric_compiled = decompress_int4_symmetric + decompress_int2_asymmetric_compiled = decompress_int2_asymmetric quantize_int8_matmul_input_compiled = quantize_int8_matmul_input unpack_int4_compiled = unpack_int4 diff --git a/modules/shared.py b/modules/shared.py index 82e290902..aafc6af11 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -518,7 +518,7 @@ options_templates.update(options_section(('quantization', "Quantization Settings "sdnq_quantize_sep": OptionInfo("

SDNQ: SDNext Quantization

", "", gr.HTML), "sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}), "sdnq_quantize_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ['pre', 'post'], "visible": native}), - "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ['int8', 'uint4', "float8_e4m3fn", 'uint8', 'int4', "float8_e5m2"], "visible": native}), + "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ['int8', 'uint4', "float8_e4m3fn", 'uint8', 'int4', "float8_e5m2", "uint2"], "visible": native}), "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": native}), "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize the convolutional layers", gr.Checkbox, {"visible": native}), "sdnq_decompress_fp32": OptionInfo(False, "Decompress using full precision", gr.Checkbox, {"visible": native}),