diff --git a/CHANGELOG.md b/CHANGELOG.md index 63b86dbd5..53b02c6b4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +16,7 @@ - **SDNQ Quantization** - Add group size support for convolutional layers - Add quantized matmul support for for convolutional layers - - Add 5-bit and 3-bit quantization support + - Add 7-bit, 5-bit and 3-bit quantization support - Fix forced FP32 with tensorwise FP8 matmul - Fix PyTorch <= 2.4 compatibility with FP8 matmul - Fix VAE with conv quant diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 70addbe7a..ccfdf1227 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -9,16 +9,16 @@ torch_version = float(torch.__version__[:3]) dtype_dict = { "int8": {"min": -128, "max": 127, "num_bits": 8, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True}, - "int7": {"min": -64, "max": 63, "num_bits": 6, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, + "int7": {"min": -64, "max": 63, "num_bits": 7, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, "int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, - "int5": {"min": -16, "max": 15, "num_bits": 5, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, + "int5": {"min": -16, "max": 15, "num_bits": 5, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "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}, "int3": {"min": -4, "max": 3, "num_bits": 3, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, "int2": {"min": -2, "max": 1, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, "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}, - "uint7": {"min": 0, "max": 127, "num_bits": 6, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, + "uint7": {"min": 0, "max": 127, "num_bits": 7, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, "uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, - "uint5": {"min": 0, "max": 31, "num_bits": 5, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, + "uint5": {"min": 0, "max": 31, "num_bits": 5, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "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}, "uint3": {"min": 0, "max": 7, "num_bits": 3, "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}, @@ -31,7 +31,7 @@ dtype_dict = { dtype_dict["bool"] = dtype_dict["uint1"] use_tensorwise_fp8_matmul = torch_version < 2.5 or devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and torch_version <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9)) -quantized_matmul_dtypes = ("int8", "int6", "int5", "int4", "int3", "int2", "float8_e4m3fn", "float8_e5m2") +quantized_matmul_dtypes = ("int8", "int7", "int6", "int5", "int4", "int3", "int2", "float8_e4m3fn", "float8_e5m2") if devices.backend in {"cpu", "openvino"}: quantized_matmul_dtypes += ("float8_e4m3fnuz", "float8_e5m2fnuz") diff --git a/modules/sdnq/decompressor.py b/modules/sdnq/decompressor.py index 2c37f26b2..c85ffc017 100644 --- a/modules/sdnq/decompressor.py +++ b/modules/sdnq/decompressor.py @@ -139,12 +139,14 @@ class PackedINTSymmetricWeightsDecompressor(torch.nn.Module): decompressor_dict = { "int8": SymmetricWeightsDecompressor, + "int7": PackedINTSymmetricWeightsDecompressor, "int6": PackedINTSymmetricWeightsDecompressor, "int5": PackedINTSymmetricWeightsDecompressor, "int4": PackedINTSymmetricWeightsDecompressor, "int3": PackedINTSymmetricWeightsDecompressor, "int2": PackedINTSymmetricWeightsDecompressor, "uint8": AsymmetricWeightsDecompressor, + "uint7": PackedINTAsymmetricWeightsDecompressor, "uint6": PackedINTAsymmetricWeightsDecompressor, "uint5": PackedINTAsymmetricWeightsDecompressor, "uint4": PackedINTAsymmetricWeightsDecompressor, diff --git a/modules/sdnq/packed_int.py b/modules/sdnq/packed_int.py index 351dc7546..e4717086e 100644 --- a/modules/sdnq/packed_int.py +++ b/modules/sdnq/packed_int.py @@ -19,6 +19,25 @@ def unpack_int_symetric(packed_tensor: torch.CharTensor, shape: torch.Size, weig return result +def pack_uint7(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, 8) + packed_tensor = torch.stack( + ( + torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 1), 128)), + torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 2), 128)), + torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128)), + torch.bitwise_or(packed_tensor[:, 3], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128)), + torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128)), + torch.bitwise_or(packed_tensor[:, 5], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 6), 128)), + torch.bitwise_or(packed_tensor[:, 6], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 7), 128)), + ), + dim=-1 + ) + return packed_tensor + + def pack_uint6(tensor: torch.Tensor) -> torch.Tensor: if tensor.dtype != torch.uint8: raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.uint8 type is supported.") @@ -109,6 +128,41 @@ def pack_uint2(tensor: torch.Tensor) -> torch.Tensor: return packed_tensor +def unpack_uint7(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: + result = torch.stack( + ( + torch.bitwise_and(packed_tensor[:, 0], 127), + torch.bitwise_and(packed_tensor[:, 1], 127), + torch.bitwise_and(packed_tensor[:, 2], 127), + torch.bitwise_and(packed_tensor[:, 3], 127), + torch.bitwise_and(packed_tensor[:, 4], 127), + torch.bitwise_and(packed_tensor[:, 5], 127), + torch.bitwise_and(packed_tensor[:, 6], 127), + torch.bitwise_or( + torch.bitwise_or( + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 1), 64), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 2), 32), + ), + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 16), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 4), 8), + ), + ), + torch.bitwise_or( + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 5), 4), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 5], 6), 2), + ), + torch.bitwise_right_shift(packed_tensor[:, 6], 7), + ), + ) + ), + dim=-1 + ).reshape(shape) + return result + + def unpack_uint6(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: result = torch.stack( ( @@ -199,11 +253,13 @@ def unpack_uint2(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor packed_int_function_dict = { + "int7": {"pack": pack_uint7, "unpack": unpack_uint7}, "int6": {"pack": pack_uint6, "unpack": unpack_uint6}, "int5": {"pack": pack_uint5, "unpack": unpack_uint5}, "int4": {"pack": pack_uint4, "unpack": unpack_uint4}, "int3": {"pack": pack_uint3, "unpack": unpack_uint3}, "int2": {"pack": pack_uint2, "unpack": unpack_uint2}, + "uint7": {"pack": pack_uint7, "unpack": unpack_uint7}, "uint6": {"pack": pack_uint6, "unpack": unpack_uint6}, "uint5": {"pack": pack_uint5, "unpack": unpack_uint5}, "uint4": {"pack": pack_uint4, "unpack": unpack_uint4}, diff --git a/modules/shared.py b/modules/shared.py index c05039363..0e7099c38 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: SD.Next 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", "float8_e4m3fn", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "visible": native}), + "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "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_compile": OptionInfo(devices.has_triton(), "Decompress using torch.compile", gr.Checkbox, {"visible": native}), diff --git a/wiki b/wiki index f01d19390..88b282df8 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit f01d19390603551f4713015f7842cae21e70f6f0 +Subproject commit 88b282df8d42ff396d5e925b5d167ffd6d791d7e