Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 01:04:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #3961 from vladmandic/dev

Browse Source 
SDNQ group size support for conv layers
This commit is contained in:
Disty0
2025-06-03 18:14:48 +03:00
committed by GitHub
parent 3ca9f29c8f 1c508fbda4
commit 894e15d3dc
2 changed files with 56 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+1
View File
@@ -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
modules/model_quant_sdnq.py
+55 -30
View File
@@ -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)