diff --git a/CHANGELOG.md b/CHANGELOG.md index 5816d5f9b..968bb8e1c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ Take a look at [Docs](https://github.com/vladmandic/sdnext/wiki/Docs), [Hints](h - `INT4` -> `uint4` - Add `float8_e4m3fn`, `float8_e5m2`, `float8_e4m3fnuz`, `float8_e5m2fnuz`, `int6`, `uint6`, `int2`, `uint2` and `uint1` support - Add quantized matmul support for `float8_e4m3fn` and `float8_e5m2` + - Add group size support for convolutional layers - Set the default quant mode to `pre` - Use per token input quant with int8 and fp8 quantized matmul - Implement better layer hijacks diff --git a/modules/model_quant_sdnq.py b/modules/model_quant_sdnq.py index 8231a3828..1281b05e4 100644 --- a/modules/model_quant_sdnq.py +++ b/modules/model_quant_sdnq.py @@ -11,6 +11,7 @@ from diffusers.utils import get_module_from_name from accelerate.utils import CustomDtype from modules import devices, shared +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}, @@ -56,15 +57,21 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz if layer_class_name in conv_types: if not quant_conv: return layer - reduction_axes = [i for i in range(layer.weight.ndim) if i != 0] - use_quantized_matmul = False is_conv_type = True + reduction_axes = 1 + output_channel_size, channel_size = layer.weight.shape[:2] + use_quantized_matmul = False + if dtype_dict[weights_dtype]["num_bits"] < 4: + weights_dtype = "uint4" elif layer_class_name in conv_transpose_types: if not quant_conv: return layer - reduction_axes = [i for i in range(layer.weight.ndim) if i != 1] - use_quantized_matmul = False is_conv_transpose_type = True + reduction_axes = 0 + channel_size, output_channel_size = layer.weight.shape[:2] + use_quantized_matmul = False + if dtype_dict[weights_dtype]["num_bits"] < 4: + weights_dtype = "uint4" else: is_linear_type = True reduction_axes = -1 @@ -73,35 +80,53 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and channel_size >= 32 and output_channel_size >= 32 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 - use_tensorwise_fp8_matmul = devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and float(torch.__version__[:3]) <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9)) + 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)) - 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 - else: - group_size = 64 - num_of_groups = channel_size // group_size + if group_size == 0: + if is_linear_type: + if dtype_dict[weights_dtype]["num_bits"] < 6: + group_size = 2 ** (2 + dtype_dict[weights_dtype]["num_bits"]) + else: + if dtype_dict[weights_dtype]["num_bits"] < 8: + group_size = 2 ** (1 + dtype_dict[weights_dtype]["num_bits"]) - if group_size >= channel_size: - group_size = channel_size - num_of_groups = 1 - else: - num_of_groups = channel_size // group_size - while channel_size % group_size != 0: # find something divisible - num_of_groups -= 1 - if num_of_groups <= 1: - group_size = channel_size - num_of_groups = 1 - break - group_size = channel_size / num_of_groups + if not use_quantized_matmul and group_size > 0: + if group_size >= channel_size: + group_size = channel_size + num_of_groups = 1 + else: + num_of_groups = channel_size // group_size + while channel_size % group_size != 0: # find something divisible + num_of_groups -= 1 + if num_of_groups <= 1: + group_size = channel_size + num_of_groups = 1 + break + group_size = channel_size / num_of_groups + group_size = int(group_size) + num_of_groups = int(num_of_groups) - if num_of_groups > 1: - result_shape = layer.weight.shape - new_shape = list(result_shape) + if num_of_groups > 1: + result_shape = layer.weight.shape + new_shape = list(result_shape) + if is_conv_type: + # output_channel_size, channel_size, X, X + # output_channel_size, num_of_groups, group_size, X, X + new_shape[1] = group_size + new_shape.insert(1, num_of_groups) + reduction_axes = 2 + elif is_conv_transpose_type: + #channel_size, output_channel_size, X, X + #num_of_groups, group_size, output_channel_size, X, X + new_shape[0] = group_size + new_shape.insert(0, num_of_groups) + reduction_axes = 1 + elif is_linear_type: + # output_channel_size, channel_size + # output_channel_size, num_of_groups, group_size last_dim_index = layer.weight.ndim - new_shape[last_dim_index - 1 : last_dim_index] = (int(num_of_groups), int(group_size)) - layer.weight.data = layer.weight.reshape(new_shape) + new_shape[last_dim_index - 1 : last_dim_index] = (num_of_groups, group_size) + layer.weight.data = layer.weight.reshape(new_shape) layer.weight.requires_grad = False if shared.opts.diffusers_offload_mode in {"none", "model"}: @@ -128,7 +153,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz zero_point = None layer.weight.data = quantize_weight(layer.weight, scale, zero_point, weights_dtype) - if not shared.opts.sdnq_decompress_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]): + if not shared.opts.sdnq_decompress_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"] and not use_tensorwise_fp8_matmul): scale = scale.to(torch_dtype) if zero_point is not None: zero_point = zero_point.to(torch_dtype)