SDNQ add 7 bit support

This commit is contained in:
Disty0
2025-06-10 11:33:06 +03:00
parent 5bd7a08877
commit 33fadf946b
6 changed files with 66 additions and 8 deletions
+1 -1
View File
@@ -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
+5 -5
View File
@@ -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")
+2
View File
@@ -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,
+56
View File
@@ -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},
+1 -1
View File
@@ -518,7 +518,7 @@ options_templates.update(options_section(("quantization", "Quantization Settings
"sdnq_quantize_sep": OptionInfo("<h2>SDNQ: SD.Next Quantization</h2>", "", 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}),
+1 -1
Submodule wiki updated: f01d193906...88b282df8d