SDNQ add 6-bit support

This commit is contained in:
Disty0
2025-05-30 12:20:13 +03:00
parent 98a11fc86c
commit 29bd2af779
3 changed files with 86 additions and 41 deletions
+72 -27
View File
@@ -18,6 +18,8 @@ debug = os.environ.get('SD_QUANT_DEBUG', None) is not None
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},
"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},
"int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
"uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": CustomDtype.INT4, "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},
@@ -30,7 +32,7 @@ dtype_dict = {
if hasattr(torch, "float8_e8m0fnu"):
dtype_dict["float8_e8m0fnu"] = {"min": 5.87747e-39, "max": 1.70141e+38, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e8m0fnu, "storage_dtype": torch.float8_e8m0fnu, "is_unsigned": True, "is_integer": False}
quantized_matmul_dtypes = ("int8", "int4", "float8_e4m3fn")
quantized_matmul_dtypes = ("int8", "int6", "int4", "float8_e4m3fn")
linear_types = ("Linear",)
conv_types = ("Conv1d", "Conv2d", "Conv3d")
@@ -73,7 +75,7 @@ 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"] < 8 and group_size != -1)):
if not use_quantized_matmul and (group_size > 0 or (dtype_dict[weights_dtype]["num_bits"] < 6 and group_size != -1)):
if group_size == 0:
if dtype_dict[weights_dtype]["num_bits"] < 4:
group_size = 32
@@ -256,11 +258,32 @@ def decompress_packed_int_asymmetric(input: torch.Tensor, scale: torch.Tensor, z
return decompress_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape), scale, zero_point, dtype, result_shape)
def decompress_int4_symmetric(input: torch.Tensor, scale: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.Tensor:
def decompress_packed_int_symmetric(input: torch.Tensor, scale: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.Tensor:
if skip_quantized_matmul:
return decompress_symmetric(unpack_int4(input, shape, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
return decompress_symmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
else:
return decompress_symmetric(unpack_int4(input, shape, dtype=scale.dtype), scale, dtype, result_shape)
return decompress_symmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape, dtype=scale.dtype), scale, dtype, result_shape)
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.")
packed_tensor = tensor.contiguous().reshape(-1, 4)
packed_tensor = torch.stack(
(
torch.bitwise_or(torch.bitwise_and(packed_tensor[:, 0], 63), torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 2), 192)),
torch.bitwise_or(torch.bitwise_and(packed_tensor[:, 1], 63), torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 4), 192)),
torch.bitwise_or(torch.bitwise_and(packed_tensor[:, 2], 63), torch.bitwise_left_shift(packed_tensor[:, 3], 6)),
),
dim=-1
)
return packed_tensor
def pack_int6(tensor: torch.Tensor) -> torch.Tensor:
if tensor.dtype != torch.int8:
raise RuntimeError(f"Invalid tensor dtype {tensor.type}. torch.int8 type is supported.")
return pack_uint6((tensor + 32).to(dtype=torch.uint8))
def pack_uint4(tensor: torch.Tensor) -> torch.Tensor:
@@ -323,6 +346,32 @@ def pack_uint1(tensor: torch.Tensor) -> torch.Tensor:
return packed_tensor
def unpack_uint6(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor:
result = torch.stack(
(
torch.bitwise_and(packed_tensor[:, 0], 63),
torch.bitwise_and(packed_tensor[:, 1], 63),
torch.bitwise_and(packed_tensor[:, 2], 63),
torch.bitwise_or(
torch.bitwise_or(
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 2), 48),
torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 4), 12),
),
torch.bitwise_right_shift(packed_tensor[:, 2], 6)
)
),
dim=-1
).reshape(shape)
return result
def unpack_int6(packed_tensor: torch.Tensor, shape: torch.Size, dtype: Optional[torch.dtype] = torch.int8, transpose: Optional[bool] = False) -> torch.Tensor:
result = unpack_uint6(packed_tensor, shape).to(dtype=dtype).sub_(32)
if transpose:
result = result.transpose(0,1)
return result
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
@@ -430,13 +479,14 @@ def int8_matmul(
bias: torch.FloatTensor,
scale: torch.FloatTensor,
compressed_weight_shape: torch.Size,
weights_dtype: str,
) -> torch.FloatTensor:
if compressed_weight_shape is not None:
weight = unpack_int4_compiled(weight, compressed_weight_shape, transpose=True)
weight = packed_int_function_dict[weights_dtype]["unpack"](weight, compressed_weight_shape, transpose=True)
return_dtype = input.dtype
output_shape = list(input.shape)
output_shape[-1] = weight.shape[-1]
input, scale = quantize_int8_matmul_input_compiled(input, scale)
input, scale = quantize_int8_matmul_input(input, scale)
result = decompress_symmetric_compiled(torch._int_mm(input, weight), scale, return_dtype, output_shape)
if bias is not None:
result.add_(bias)
@@ -454,7 +504,7 @@ def quantized_linear_forward_fp8_matmul_sm89(self, input: torch.FloatTensor) ->
def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor:
if torch.numel(input) / input.shape[-1] < 32:
return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight, skip_quantized_matmul=True), self.bias)
return int8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale, getattr(self.sdnq_decompressor, "compressed_weight_shape", None))
return int8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale, getattr(self.sdnq_decompressor, "compressed_weight_shape", None), self.sdnq_decompressor.weights_dtype)
def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
@@ -556,18 +606,19 @@ class PackedINTAsymmetricWeightsDecompressor(torch.nn.Module):
return decompress_packed_int_asymmetric_compiled(weight, self.scale, self.zero_point, self.compressed_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype)
class INT4SymmetricWeightsDecompressor(torch.nn.Module):
class PackedINTSymmetricWeightsDecompressor(torch.nn.Module):
def __init__(
self,
scale: torch.Tensor,
compressed_weight_shape: torch.Size,
result_dtype: torch.dtype,
result_shape: torch.Size,
weights_dtype: str,
use_quantized_matmul: bool = False,
**kwargs,
):
super().__init__()
self.weights_dtype = "int4"
self.weights_dtype = weights_dtype
self.use_quantized_matmul = use_quantized_matmul
self.compressed_weight_shape = compressed_weight_shape
self.result_dtype = result_dtype
@@ -575,16 +626,18 @@ class INT4SymmetricWeightsDecompressor(torch.nn.Module):
self.register_buffer("scale", scale)
def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
return pack_int4(weight.to(dtype=torch.int8))
return packed_int_function_dict[self.weights_dtype]["pack"](weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]))
def forward(self, weight, skip_quantized_matmul=False, **kwargs):
return decompress_int4_symmetric_compiled(weight, self.scale, self.compressed_weight_shape, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul)
return decompress_packed_int_symmetric_compiled(weight, self.scale, self.compressed_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul)
decompressor_dict = {
"int8": SymmetricWeightsDecompressor,
"uint8": AsymmetricWeightsDecompressor,
"int4": INT4SymmetricWeightsDecompressor,
"int6": PackedINTSymmetricWeightsDecompressor,
"uint6": PackedINTAsymmetricWeightsDecompressor,
"int4": PackedINTSymmetricWeightsDecompressor,
"uint4": PackedINTAsymmetricWeightsDecompressor,
"uint2": PackedINTAsymmetricWeightsDecompressor,
"uint1": PackedINTAsymmetricWeightsDecompressor,
@@ -597,6 +650,8 @@ decompressor_dict = {
packed_int_function_dict = {
"int6": {"pack": pack_int6, "unpack": unpack_int6},
"uint6": {"pack": pack_uint6, "unpack": unpack_uint6},
"int4": {"pack": pack_int4, "unpack": unpack_int4},
"uint4": {"pack": pack_uint4, "unpack": unpack_uint4},
"uint2": {"pack": pack_uint2, "unpack": unpack_uint2},
@@ -807,28 +862,18 @@ if shared.opts.sdnq_decompress_compile:
decompress_asymmetric_compiled = torch.compile(decompress_asymmetric, fullgraph=True)
decompress_symmetric_compiled = torch.compile(decompress_symmetric, fullgraph=True)
decompress_packed_int_asymmetric_compiled = torch.compile(decompress_packed_int_asymmetric, fullgraph=True)
decompress_int4_symmetric_compiled = torch.compile(decompress_int4_symmetric, fullgraph=True)
decompress_packed_int_symmetric_compiled = torch.compile(decompress_packed_int_symmetric, 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
quantize_int8_matmul_input_compiled = quantize_int8_matmul_input
unpack_int4_compiled = unpack_int4
int8_matmul = torch.compile(int8_matmul, fullgraph=True)
else:
quantize_int8_matmul_input_compiled = torch.compile(quantize_int8_matmul_input, fullgraph=True)
unpack_int4_compiled = torch.compile(unpack_int4, fullgraph=True)
int8_matmul = torch.compile(int8_matmul, fullgraph=True)
except Exception as e:
shared.log.warning(f"Quantization: type=sdnq Decompress using torch.compile is not available: {e}")
decompress_asymmetric_compiled = decompress_asymmetric
decompress_symmetric_compiled = decompress_symmetric
decompress_packed_int_asymmetric_compiled = decompress_packed_int_asymmetric
decompress_int4_symmetric_compiled = decompress_int4_symmetric
quantize_int8_matmul_input_compiled = quantize_int8_matmul_input
unpack_int4_compiled = unpack_int4
decompress_packed_int_symmetric_compiled = decompress_packed_int_symmetric
else:
decompress_asymmetric_compiled = decompress_asymmetric
decompress_symmetric_compiled = decompress_symmetric
decompress_packed_int_asymmetric_compiled = decompress_packed_int_asymmetric
decompress_int4_symmetric_compiled = decompress_int4_symmetric
quantize_int8_matmul_input_compiled = quantize_int8_matmul_input
unpack_int4_compiled = unpack_int4
decompress_packed_int_symmetric_compiled = decompress_packed_int_symmetric
+12 -12
View File
@@ -505,20 +505,20 @@ options_templates.update(options_section(('backends', "Backend Settings"), {
"openvino_sep": OptionInfo("<h2>OpenVINO</h2>", "", gr.HTML, {"visible": cmd_opts.use_openvino}),
"openvino_devices": OptionInfo([], "OpenVINO devices to use", gr.CheckboxGroup, {"choices": get_openvino_device_list() if cmd_opts.use_openvino else [], "visible": cmd_opts.use_openvino}), # pylint: disable=E0606
"openvino_accuracy": OptionInfo("performance", "OpenVINO accuracy mode", gr.Radio, {"choices": ['performance', 'accuracy'], "visible": cmd_opts.use_openvino}),
"openvino_accuracy": OptionInfo("performance", "OpenVINO accuracy mode", gr.Radio, {"choices": ["performance", "accuracy"], "visible": cmd_opts.use_openvino}),
"openvino_disable_model_caching": OptionInfo(True, "OpenVINO disable model caching", gr.Checkbox, {"visible": cmd_opts.use_openvino}),
"openvino_disable_memory_cleanup": OptionInfo(True, "OpenVINO disable memory cleanup after compile", gr.Checkbox, {"visible": cmd_opts.use_openvino}),
"directml_sep": OptionInfo("<h2>DirectML</h2>", "", gr.HTML, {"visible": devices.backend == "directml"}),
"directml_memory_provider": OptionInfo(default_memory_provider, 'DirectML memory stats provider', gr.Radio, {"choices": memory_providers, "visible": devices.backend == "directml"}),
"directml_memory_provider": OptionInfo(default_memory_provider, "DirectML memory stats provider", gr.Radio, {"choices": memory_providers, "visible": devices.backend == "directml"}),
"directml_catch_nan": OptionInfo(False, "DirectML retry ops for NaN", gr.Checkbox, {"visible": devices.backend == "directml"}),
}))
options_templates.update(options_section(('quantization', "Quantization Settings"), {
options_templates.update(options_section(("quantization", "Quantization Settings"), {
"sdnq_quantize_sep": OptionInfo("<h2>SDNQ: SDNext 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', 'uint4', "float8_e4m3fn", 'uint8', 'int4', "float8_e5m2", "uint2", "uint1"], "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", "int6", "uint4", "float8_e4m3fn", "uint8", "uint6", "int4", "float8_e5m2", "uint2", "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_fp32": OptionInfo(False, "Decompress using full precision", gr.Checkbox, {"visible": native}),
@@ -529,7 +529,7 @@ options_templates.update(options_section(('quantization', "Quantization Settings
"bnb_quantization_sep": OptionInfo("<h2>BitsAndBytes</h2>", "", gr.HTML),
"bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
"bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ['nf4', 'fp8', 'fp4'], "visible": native}),
"bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ["nf4", "fp8", "fp4"], "visible": native}),
"bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Dropdown, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"], "visible": native}),
"quanto_quantization_sep": OptionInfo("<h2>Optimum Quanto</h2>", "", gr.HTML),
@@ -538,14 +538,14 @@ options_templates.update(options_section(('quantization', "Quantization Settings
"optimum_quanto_sep": OptionInfo("<h2>Optimum Quanto: post-load</h2>", "", gr.HTML),
"optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
"optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}),
"optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ['none', 'qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2'], "visible": native}),
"optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ["qint8", "qfloat8_e4m3fn", "qfloat8_e5m2", "qint4", "qint2"], "visible": native}),
"optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ["none", "qint8", "qfloat8_e4m3fn", "qfloat8_e5m2"], "visible": native}),
"optimum_quanto_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox, {"visible": native}),
"torchao_sep": OptionInfo("<h2>TorchAO</h2>", "", gr.HTML),
"torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM"], "visible": native}),
"torchao_quantization_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ['pre', 'post'], "visible": native}),
"torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ['int4_weight_only', 'int8_dynamic_activation_int4_weight', 'int8_weight_only', 'int8_dynamic_activation_int8_weight', 'float8_weight_only', 'float8_dynamic_activation_float8_weight', 'float8_static_activation_float8_weight'], "visible": native}),
"torchao_quantization_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ["pre", "post"], "visible": native}),
"torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ["int4_weight_only", "int8_dynamic_activation_int4_weight", "int8_weight_only", "int8_dynamic_activation_int8_weight", "float8_weight_only", "float8_dynamic_activation_float8_weight", "float8_static_activation_float8_weight"], "visible": native}),
"layerwise_quantization_sep": OptionInfo("<h2>Layerwise Casting</h2>", "", gr.HTML),
"layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "TE"], "visible": native}),
@@ -559,11 +559,11 @@ options_templates.update(options_section(('quantization', "Quantization Settings
"nncf_compress_sep": OptionInfo("<h2>NNCF: Neural Network Compression Framework</h2>", "", gr.HTML, {"visible": cmd_opts.use_openvino}),
"nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": cmd_opts.use_openvino}),
"nncf_compress_weights_mode": OptionInfo("INT8_SYM", "Quantization type", gr.Dropdown, {"choices": ['INT8', 'INT4_ASYM', 'INT8_SYM', 'INT4_SYM', 'NF4'], "visible": cmd_opts.use_openvino}),
"nncf_compress_weights_mode": OptionInfo("INT8_SYM", "Quantization type", gr.Dropdown, {"choices": ["INT8", "INT4_ASYM", "INT8_SYM", "INT4_SYM", "NF4"], "visible": cmd_opts.use_openvino}),
"nncf_compress_weights_raito": OptionInfo(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}),
"nncf_compress_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": cmd_opts.use_openvino}),
"nncf_quantize": OptionInfo([], "Static Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE"], "visible": cmd_opts.use_openvino}),
"nncf_quantize_mode": OptionInfo("INT8", "OpenVINO activations mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}),
"nncf_quantize_mode": OptionInfo("INT8", "OpenVINO activations mode", gr.Dropdown, {"choices": ["INT8", "FP8_E4M3", "FP8_E5M2"], "visible": cmd_opts.use_openvino}),
}))
options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {