From c3d007b02c97a9f5de03807da38774f6772f8920 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 2 Aug 2025 17:36:55 +0300 Subject: [PATCH] SDNQ split forward.py into layers and cleanup --- modules/sdnq/__init__.py | 71 ++-- modules/sdnq/common.py | 13 +- modules/sdnq/dequantizer.py | 6 +- modules/sdnq/forward.py | 388 +----------------- modules/sdnq/layers/conv/conv_fp8.py | 71 ++++ .../sdnq/layers/conv/conv_fp8_tensorwise.py | 70 ++++ modules/sdnq/layers/conv/conv_int8.py | 76 ++++ modules/sdnq/layers/conv/forward.py | 93 +++++ modules/sdnq/layers/linear/forward.py | 7 + modules/sdnq/layers/linear/linear_fp8.py | 41 ++ .../layers/linear/linear_fp8_tensorwise.py | 48 +++ modules/sdnq/layers/linear/linear_int8.py | 52 +++ modules/sdnq/packed_int.py | 1 + 13 files changed, 521 insertions(+), 416 deletions(-) create mode 100644 modules/sdnq/layers/conv/conv_fp8.py create mode 100644 modules/sdnq/layers/conv/conv_fp8_tensorwise.py create mode 100644 modules/sdnq/layers/conv/conv_int8.py create mode 100644 modules/sdnq/layers/conv/forward.py create mode 100644 modules/sdnq/layers/linear/forward.py create mode 100644 modules/sdnq/layers/linear/linear_fp8.py create mode 100644 modules/sdnq/layers/linear/linear_fp8_tensorwise.py create mode 100644 modules/sdnq/layers/linear/linear_int8.py diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index 985e99a2a..e6c012e8c 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -14,6 +14,41 @@ from .dequantizer import dequantizer_dict from .forward import get_forward_func +class QuantizationMethod(str, Enum): + SDNQ = "sdnq" + + +def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]: + zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True) + scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"]) + eps = torch.finfo(scale.dtype).eps # prevent divison by 0 + scale = torch.where(torch.abs(scale) < eps, eps, scale) + if dtype_dict[weights_dtype]["min"] != 0: + zero_point.sub_(torch.mul(scale, dtype_dict[weights_dtype]["min"])) + return scale, zero_point + + +def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> torch.FloatTensor: + scale = torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).div_(dtype_dict[weights_dtype]["max"]) + eps = torch.finfo(scale.dtype).eps # prevent divison by 0 + scale = torch.where(torch.abs(scale) < eps, eps, scale) + return scale + + +def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]: + if dtype_dict[weights_dtype]["is_unsigned"]: + scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype) + quantized_weight = torch.sub(weight, zero_point).div_(scale) + else: + scale = get_scale_symmetric(weight, reduction_axes, weights_dtype) + quantized_weight = torch.div(weight, scale) + zero_point = None + if dtype_dict[weights_dtype]["is_integer"]: + quantized_weight.round_() + quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"]) + return quantized_weight, scale, zero_point + + @devices.inference_context() def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument layer_class_name = layer.__class__.__name__ @@ -153,7 +188,6 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, dtype_dict[weights_dtype]["is_integer"], use_tensorwise_fp8_matmul) layer.forward = layer.forward.__get__(layer, layer.__class__) - #devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}") return layer @@ -195,41 +229,6 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si return model -def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]: - zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True) - scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"]) - eps = torch.finfo(scale.dtype).eps # prevent divison by 0 - scale = torch.where(torch.abs(scale) < eps, eps, scale) - if dtype_dict[weights_dtype]["min"] != 0: - zero_point.sub_(torch.mul(scale, dtype_dict[weights_dtype]["min"])) - return scale, zero_point - - -def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> torch.FloatTensor: - scale = torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).div_(dtype_dict[weights_dtype]["max"]) - eps = torch.finfo(scale.dtype).eps # prevent divison by 0 - scale = torch.where(torch.abs(scale) < eps, eps, scale) - return scale - - -def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]: - if dtype_dict[weights_dtype]["is_unsigned"]: - scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype) - quantized_weight = torch.sub(weight, zero_point).div_(scale) - else: - scale = get_scale_symmetric(weight, reduction_axes, weights_dtype) - quantized_weight = torch.div(weight, scale) - zero_point = None - if dtype_dict[weights_dtype]["is_integer"]: - quantized_weight.round_() - quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"]) - return quantized_weight, scale, zero_point - - -class QuantizationMethod(str, Enum): - SDNQ = "sdnq" - - class SDNQQuantizer(DiffusersQuantizer): r""" Diffusers Quantizer for SDNQ diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 9dab864c5..07ae0671b 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -1,7 +1,7 @@ # pylint: disable=redefined-builtin,no-member,protected-access import torch -from modules import devices +from modules import devices, shared torch_version = float(torch.__version__[:3]) @@ -30,7 +30,9 @@ if hasattr(torch, "float8_e4m3fnuz"): if hasattr(torch, "float8_e5m2fnuz"): dtype_dict["float8_e5m2fnuz"] = {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False} -use_tensorwise_fp8_matmul = True # Direct tensorwise only exist on H100 hardware, sdnq will use software tensorwise with this setting +use_torch_compile = shared.opts.sdnq_dequantize_compile # this setting requires a full restart of the webui to apply +use_tensorwise_fp8_matmul = True # row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting + 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") @@ -39,3 +41,10 @@ linear_types = ("Linear",) conv_types = ("Conv1d", "Conv2d", "Conv3d") conv_transpose_types = ("ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d") allowed_types = linear_types + conv_types + conv_transpose_types + +if use_torch_compile: + try: + torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) + torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit) + except Exception as e: + shared.log.warning(f"Quantization: type=sdnq Failed to increase the cache size for torch.compile: {e}") diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index db6fb622c..0a0881117 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -3,7 +3,7 @@ import torch from modules import shared -from .common import dtype_dict +from .common import dtype_dict, use_torch_compile from .packed_int import pack_int_symetric, unpack_int_symetric, packed_int_function_dict @@ -173,10 +173,8 @@ dequantizer_dict = { } -if shared.opts.sdnq_dequantize_compile: +if use_torch_compile: try: - torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) - torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit) dequantize_asymmetric_compiled = torch.compile(dequantize_asymmetric, fullgraph=True, dynamic=False) dequantize_symmetric_compiled = torch.compile(dequantize_symmetric, fullgraph=True, dynamic=False) dequantize_packed_int_asymmetric_compiled = torch.compile(dequantize_packed_int_asymmetric, fullgraph=True, dynamic=False) diff --git a/modules/sdnq/forward.py b/modules/sdnq/forward.py index 4abc4d7e8..211b57768 100644 --- a/modules/sdnq/forward.py +++ b/modules/sdnq/forward.py @@ -1,410 +1,50 @@ -# pylint: disable=redefined-builtin,no-member,protected-access +# pylint: disable=protected-access + +from typing import Callable -from typing import Callable, List, Tuple, Optional import torch -from modules import shared from .common import conv_types, conv_transpose_types -from .dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias -from .packed_int import unpack_int_symetric def get_forward_func(layer_class_name: str, use_quantized_matmul: bool, is_integer: bool, use_tensorwise_fp8_matmul: bool) -> Callable: # pylint: disable=inconsistent-return-statements if layer_class_name in conv_types: if use_quantized_matmul: if is_integer: + from .layers.conv.conv_int8 import quantized_conv_forward_int8_matmul return quantized_conv_forward_int8_matmul else: if use_tensorwise_fp8_matmul: + from .layers.conv.conv_fp8_tensorwise import quantized_conv_forward_fp8_matmul_tensorwise return quantized_conv_forward_fp8_matmul_tensorwise else: + from .layers.conv.conv_fp8 import quantized_conv_forward_fp8_matmul return quantized_conv_forward_fp8_matmul else: + from .layers.conv.forward import quantized_conv_forward return quantized_conv_forward elif layer_class_name in conv_transpose_types: if layer_class_name.endswith("1d"): + from .layers.conv.forward import quantized_conv_transpose_1d_forward return quantized_conv_transpose_1d_forward elif layer_class_name.endswith("2d"): + from .layers.conv.forward import quantized_conv_transpose_2d_forward return quantized_conv_transpose_2d_forward elif layer_class_name.endswith("3d"): + from .layers.conv.forward import quantized_conv_transpose_3d_forward return quantized_conv_transpose_3d_forward else: if use_quantized_matmul: if is_integer: + from .layers.linear.linear_int8 import quantized_linear_forward_int8_matmul return quantized_linear_forward_int8_matmul else: if use_tensorwise_fp8_matmul: + from .layers.linear.linear_fp8_tensorwise import quantized_linear_forward_fp8_matmul_tensorwise return quantized_linear_forward_fp8_matmul_tensorwise else: + from .layers.linear.linear_fp8 import quantized_linear_forward_fp8_matmul return quantized_linear_forward_fp8_matmul else: + from .layers.linear.forward import quantized_linear_forward return quantized_linear_forward - - -def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]: - input = input.flatten(0,-2).contiguous() - input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448) - input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn) - input_scale = input_scale.to(dtype=torch.float32) - return input, input_scale - - -def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]: - input = input.flatten(0,-2).contiguous() - input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448) - input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn) - scale = torch.mul(input_scale, scale) - if scale.dtype == torch.float16: # fp16 will overflow - scale = scale.to(dtype=torch.float32) - return input, scale - - -def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]: - input = input.flatten(0,-2).contiguous() - input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(127) - input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(dtype=torch.int8) - scale = torch.mul(input_scale, scale) - if scale.dtype == torch.float16: # fp16 will overflow - scale = scale.to(dtype=torch.float32) - return input, scale - - -def fp8_matmul( - input: torch.FloatTensor, - weight: torch.Tensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, -) -> torch.FloatTensor: - return_dtype = input.dtype - output_shape = list(input.shape) - output_shape[-1] = weight.shape[-1] - input, input_scale = quantize_fp8_matmul_input(input) - return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(output_shape) - - -# sm89 doesn't support row wise scale in Windows -def fp8_matmul_tensorwise( - input: torch.FloatTensor, - weight: torch.Tensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, -) -> torch.FloatTensor: - return_dtype = input.dtype - output_shape = list(input.shape) - output_shape[-1] = weight.shape[-1] - dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) - input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) - if bias is not None: - return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, bias, return_dtype, output_shape) - else: - return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape) - - -def int8_matmul( - input: torch.FloatTensor, - weight: torch.Tensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - quantized_weight_shape: torch.Size, - weights_dtype: str, -) -> torch.FloatTensor: - if quantized_weight_shape is not None: - weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True) - return_dtype = input.dtype - output_shape = list(input.shape) - output_shape[-1] = weight.shape[-1] - input, scale = quantize_int8_matmul_input(input, scale) - if bias is not None: - return dequantize_symmetric_with_bias(torch._int_mm(input, weight), scale, bias, return_dtype, output_shape) - else: - return dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape) - - -def process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation): - if conv_type == 1: - batch_size, _, L_in = input.shape - C_out, _, K_l = result_shape - L_out = (L_in + 2 * padding[1] - dilation[1] * (K_l - 1) - 1) // stride[1] + 1 - mm_output_shape = (batch_size, L_out, C_out) - kernel_size = (1, K_l) - if conv_type == 2: - batch_size, _, H_in, W_in = input.shape - C_out, _, K_h, K_w = result_shape - H_out = (H_in + 2 * padding[0] - dilation[0] * (K_h - 1) - 1) // stride[0] + 1 - W_out = (W_in + 2 * padding[1] - dilation[1] * (K_w - 1) - 1) // stride[1] + 1 - mm_output_shape = (batch_size, H_out, W_out, C_out) - kernel_size = (K_h, K_w) - else: - batch_size, _, D_in, H_in, W_in = input.shape - C_out, _, K_d, K_h, K_w = result_shape - D_out = (D_in + 2 * padding[0] - dilation[0] * (K_d - 1) - 1) // stride[0] + 1 - H_out = (H_in + 2 * padding[1] - dilation[1] * (K_h - 1) - 1) // stride[1] + 1 - W_out = (W_in + 2 * padding[2] - dilation[2] * (K_w - 1) - 1) // stride[2] + 1 - mm_output_shape = (batch_size, D_out, H_out, W_out, C_out) - kernel_size = (K_d, K_h, K_w) - - if padding_mode != "zeros": - input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode) - padding = (0,) * (conv_type if conv_type != 1 else 2) - elif conv_type == 3: - input = torch.nn.functional.pad(input, reversed_padding_repeated_twice) - - if conv_type == 1: - input = input.unsqueeze(2) - - if conv_type == 3: - K_D_eff = K_d + (K_d - 1) * (dilation[0] - 1) - K_H_eff = K_h + (K_h - 1) * (dilation[0] - 1) - K_W_eff = K_w + (K_w - 1) * (dilation[0] - 1) - input = input.unfold(2, K_D_eff, stride[0]).unfold(3, K_H_eff, stride[1]).unfold(4, K_W_eff, stride[2]) - if dilation[0] > 1: - input = input[..., ::dilation[0], :, :] - if dilation[1] > 1: - input = input[..., ::dilation[1], :] - if dilation[2] > 1: - input = input[..., ::dilation[2]] - input = input.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(batch_size, D_out * H_out * W_out, -1) - else: - input = torch.nn.functional.unfold(input, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation).transpose(1,2) - return input, mm_output_shape - - -def conv_fp8_matmul( - input: torch.FloatTensor, - weight: torch.Tensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - result_shape: torch.Size, - reversed_padding_repeated_twice: List[int], - padding_mode: str, conv_type: int, - groups: int, stride: List[int], - padding: List[int], dilation: List[int], -) -> torch.FloatTensor: - return_dtype = input.dtype - input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) - input, input_scale = quantize_fp8_matmul_input(input) - - if groups == 1: - result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(mm_output_shape) - else: - scale = scale.reshape(groups, 1, scale.shape[1] // groups) - input_scale = input_scale.reshape(groups, input_scale.shape[0] // groups, 1) - weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) - input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) - result = [] - if bias is not None: - bias = bias.reshape(groups, bias.shape[0] // groups) - for i in range(groups): - result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype)) - else: - for i in range(groups): - result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype)) - result = torch.cat(result, dim=-1).reshape(mm_output_shape) - - if conv_type == 1: - result = result.transpose(1,2) - elif conv_type == 2: - result = result.permute(0,3,1,2) - elif conv_type == 3: - result = result.permute(0,4,1,2,3) - return result - - -def conv_fp8_matmul_tensorwise( - input: torch.FloatTensor, - weight: torch.Tensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - result_shape: torch.Size, - reversed_padding_repeated_twice: List[int], - padding_mode: str, conv_type: int, - groups: int, stride: List[int], - padding: List[int], dilation: List[int], -) -> torch.FloatTensor: - return_dtype = input.dtype - input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) - input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) - dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) - - if groups == 1: - result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype) - else: - weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) - input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) - result = [] - for i in range(groups): - result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype)) - result = torch.cat(result, dim=-1) - if bias is not None: - dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape) - else: - dequantize_symmetric(result, scale, return_dtype, mm_output_shape) - - if conv_type == 1: - result = result.transpose(1,2) - elif conv_type == 2: - result = result.permute(0,3,1,2) - elif conv_type == 3: - result = result.permute(0,4,1,2,3) - return result - - -def conv_int8_matmul( - input: torch.FloatTensor, - weight: torch.CharTensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - result_shape: torch.Size, - quantized_weight_shape: torch.Size, - weights_dtype: str, - reversed_padding_repeated_twice: List[int], - padding_mode: str, conv_type: int, - groups: int, stride: List[int], - padding: List[int], dilation: List[int], -) -> torch.FloatTensor: - return_dtype = input.dtype - input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) - input, scale = quantize_int8_matmul_input(input, scale) - if quantized_weight_shape is not None: - weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True) - - if groups == 1: - result = torch._int_mm(input, weight) - else: - weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) - input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) - result = [] - for i in range(groups): - result.append(torch._int_mm(input[i], weight[i])) - result = torch.cat(result, dim=-1) - if bias is not None: - result = dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape) - else: - result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape) - - if conv_type == 1: - result = result.transpose(1,2) - elif conv_type == 2: - result = result.permute(0,3,1,2) - elif conv_type == 3: - result = result.permute(0,4,1,2,3) - return result - - -def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: - if torch.numel(input) / input.shape[-1] < 32: - return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - return fp8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale) - - -def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor: - if torch.numel(input) / input.shape[-1] < 32: - return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_dequantizer.scale) - - -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_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - return int8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale, getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), self.sdnq_dequantizer.weights_dtype) - - -def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor: - return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias) - - -def get_conv_args(input_ndim: int, stride, padding, dilation): - if input_ndim == 3: - conv_type = 1 - elif input_ndim == 4: - conv_type = 2 - else: - conv_type = 3 - if isinstance(stride, int): - stride = (stride,) * conv_type - if isinstance(padding, int): - padding = (padding,) * conv_type - if isinstance(dilation, int): - dilation = (dilation,) * conv_type - if conv_type == 1: - stride = (1, stride[0]) - padding = (0, padding[0]) - dilation = (1, dilation[0]) - return conv_type, stride, padding, dilation - - -def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor: - if torch.numel(input) / input.shape[2] < 32: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) - return conv_fp8_matmul( - input, self.weight, self.bias, - self.sdnq_dequantizer.scale, - self.sdnq_dequantizer.result_shape, - self._reversed_padding_repeated_twice, - self.padding_mode, conv_type, - self.groups, stride, padding, dilation, - ) - - -def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTensor: - if torch.numel(input) / input.shape[2] < 32: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) - return conv_fp8_matmul_tensorwise( - input, self.weight, self.bias, - self.sdnq_dequantizer.scale, - self.sdnq_dequantizer.result_shape, - self._reversed_padding_repeated_twice, - self.padding_mode, conv_type, - self.groups, stride, padding, dilation, - ) - - -def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor: - if torch.numel(input) / input.shape[2] < 32: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) - return conv_int8_matmul( - input, self.weight, self.bias, - self.sdnq_dequantizer.scale, - self.sdnq_dequantizer.result_shape, - getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), - self.sdnq_dequantizer.weights_dtype, - self._reversed_padding_repeated_twice, - self.padding_mode, conv_type, - self.groups, stride, padding, dilation, - ) - - -def quantized_conv_forward(self, input) -> torch.FloatTensor: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight), self.bias) - - -def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: - output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation) - return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) - - -def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: - output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation) - return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) - - -def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: - output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation) - return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) - - -if shared.opts.sdnq_dequantize_compile: - try: - torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) - torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit) - int8_matmul = torch.compile(int8_matmul, fullgraph=True, dynamic=False) - fp8_matmul = torch.compile(fp8_matmul, fullgraph=True, dynamic=False) - fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True, dynamic=False) - conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True, dynamic=False) - conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True, dynamic=False) - conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True, dynamic=False) - except Exception as e: - shared.log.warning(f"Quantization: type=sdnq MatMul using torch.compile is not available: {e}") diff --git a/modules/sdnq/layers/conv/conv_fp8.py b/modules/sdnq/layers/conv/conv_fp8.py new file mode 100644 index 000000000..1f5a50922 --- /dev/null +++ b/modules/sdnq/layers/conv/conv_fp8.py @@ -0,0 +1,71 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import List + +import torch + +from ...common import use_torch_compile # noqa: TID252 +from ..linear.linear_fp8 import quantize_fp8_matmul_input # noqa: TID252 +from .conv import get_conv_args, process_conv_input + + +def conv_fp8_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, + result_shape: torch.Size, + reversed_padding_repeated_twice: List[int], + padding_mode: str, conv_type: int, + groups: int, stride: List[int], + padding: List[int], dilation: List[int], +) -> torch.FloatTensor: + return_dtype = input.dtype + input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) + input, input_scale = quantize_fp8_matmul_input(input) + + if groups == 1: + result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(mm_output_shape) + else: + scale = scale.reshape(groups, 1, scale.shape[1] // groups) + input_scale = input_scale.reshape(groups, input_scale.shape[0] // groups, 1) + weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) + input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) + result = [] + if bias is not None: + bias = bias.reshape(groups, bias.shape[0] // groups) + for i in range(groups): + result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype)) + else: + for i in range(groups): + result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype)) + result = torch.cat(result, dim=-1).reshape(mm_output_shape) + + if conv_type == 1: + result = result.transpose(1,2) + elif conv_type == 2: + result = result.permute(0,3,1,2) + elif conv_type == 3: + result = result.permute(0,4,1,2,3) + return result + + +def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor: + if torch.numel(input) / input.shape[2] < 32: + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) + return conv_fp8_matmul( + input, self.weight, self.bias, + self.sdnq_dequantizer.scale, + self.sdnq_dequantizer.result_shape, + self._reversed_padding_repeated_twice, + self.padding_mode, conv_type, + self.groups, stride, padding, dilation, + ) + + +if use_torch_compile: + try: + conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True, dynamic=False) + except Exception: + pass diff --git a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py new file mode 100644 index 000000000..3c5cfa56d --- /dev/null +++ b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py @@ -0,0 +1,70 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import List + +import torch + +from ...common import use_torch_compile # noqa: TID252 +from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 +from ..linear.linear_fp8_tensorwise import quantize_fp8_matmul_input_tensorwise # noqa: TID252 +from .conv import get_conv_args, process_conv_input + + +def conv_fp8_matmul_tensorwise( + input: torch.FloatTensor, + weight: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, + result_shape: torch.Size, + reversed_padding_repeated_twice: List[int], + padding_mode: str, conv_type: int, + groups: int, stride: List[int], + padding: List[int], dilation: List[int], +) -> torch.FloatTensor: + return_dtype = input.dtype + input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) + input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) + dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) + + if groups == 1: + result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype) + else: + weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) + input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) + result = [] + for i in range(groups): + result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype)) + result = torch.cat(result, dim=-1) + if bias is not None: + dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape) + else: + dequantize_symmetric(result, scale, return_dtype, mm_output_shape) + + if conv_type == 1: + result = result.transpose(1,2) + elif conv_type == 2: + result = result.permute(0,3,1,2) + elif conv_type == 3: + result = result.permute(0,4,1,2,3) + return result + + +def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTensor: + if torch.numel(input) / input.shape[2] < 32: + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) + return conv_fp8_matmul_tensorwise( + input, self.weight, self.bias, + self.sdnq_dequantizer.scale, + self.sdnq_dequantizer.result_shape, + self._reversed_padding_repeated_twice, + self.padding_mode, conv_type, + self.groups, stride, padding, dilation, + ) + + +if use_torch_compile: + try: + conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True, dynamic=False) + except Exception: + pass diff --git a/modules/sdnq/layers/conv/conv_int8.py b/modules/sdnq/layers/conv/conv_int8.py new file mode 100644 index 000000000..ffda82e2f --- /dev/null +++ b/modules/sdnq/layers/conv/conv_int8.py @@ -0,0 +1,76 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import List + +import torch + +from ...common import use_torch_compile # noqa: TID252 +from ...packed_int import unpack_int_symetric # noqa: TID252 +from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 +from ..linear.linear_int8 import quantize_int8_matmul_input # noqa: TID252 +from .conv import get_conv_args, process_conv_input + + +def conv_int8_matmul( + input: torch.FloatTensor, + weight: torch.CharTensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, + result_shape: torch.Size, + quantized_weight_shape: torch.Size, + weights_dtype: str, + reversed_padding_repeated_twice: List[int], + padding_mode: str, conv_type: int, + groups: int, stride: List[int], + padding: List[int], dilation: List[int], +) -> torch.FloatTensor: + return_dtype = input.dtype + input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) + input, scale = quantize_int8_matmul_input(input, scale) + if quantized_weight_shape is not None: + weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True) + + if groups == 1: + result = torch._int_mm(input, weight) + else: + weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) + input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) + result = [] + for i in range(groups): + result.append(torch._int_mm(input[i], weight[i])) + result = torch.cat(result, dim=-1) + if bias is not None: + result = dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape) + else: + result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape) + + if conv_type == 1: + result = result.transpose(1,2) + elif conv_type == 2: + result = result.permute(0,3,1,2) + elif conv_type == 3: + result = result.permute(0,4,1,2,3) + return result + + +def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor: + if torch.numel(input) / input.shape[2] < 32: + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) + return conv_int8_matmul( + input, self.weight, self.bias, + self.sdnq_dequantizer.scale, + self.sdnq_dequantizer.result_shape, + getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), + self.sdnq_dequantizer.weights_dtype, + self._reversed_padding_repeated_twice, + self.padding_mode, conv_type, + self.groups, stride, padding, dilation, + ) + + +if use_torch_compile: + try: + conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True, dynamic=False) + except Exception: + pass diff --git a/modules/sdnq/layers/conv/forward.py b/modules/sdnq/layers/conv/forward.py new file mode 100644 index 000000000..44c90b061 --- /dev/null +++ b/modules/sdnq/layers/conv/forward.py @@ -0,0 +1,93 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import Optional + +import torch + + +def get_conv_args(input_ndim: int, stride, padding, dilation): + if input_ndim == 3: + conv_type = 1 + elif input_ndim == 4: + conv_type = 2 + else: + conv_type = 3 + if isinstance(stride, int): + stride = (stride,) * conv_type + if isinstance(padding, int): + padding = (padding,) * conv_type + if isinstance(dilation, int): + dilation = (dilation,) * conv_type + if conv_type == 1: + stride = (1, stride[0]) + padding = (0, padding[0]) + dilation = (1, dilation[0]) + return conv_type, stride, padding, dilation + + +def process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation): + if conv_type == 1: + batch_size, _, L_in = input.shape + C_out, _, K_l = result_shape + L_out = (L_in + 2 * padding[1] - dilation[1] * (K_l - 1) - 1) // stride[1] + 1 + mm_output_shape = (batch_size, L_out, C_out) + kernel_size = (1, K_l) + if conv_type == 2: + batch_size, _, H_in, W_in = input.shape + C_out, _, K_h, K_w = result_shape + H_out = (H_in + 2 * padding[0] - dilation[0] * (K_h - 1) - 1) // stride[0] + 1 + W_out = (W_in + 2 * padding[1] - dilation[1] * (K_w - 1) - 1) // stride[1] + 1 + mm_output_shape = (batch_size, H_out, W_out, C_out) + kernel_size = (K_h, K_w) + else: + batch_size, _, D_in, H_in, W_in = input.shape + C_out, _, K_d, K_h, K_w = result_shape + D_out = (D_in + 2 * padding[0] - dilation[0] * (K_d - 1) - 1) // stride[0] + 1 + H_out = (H_in + 2 * padding[1] - dilation[1] * (K_h - 1) - 1) // stride[1] + 1 + W_out = (W_in + 2 * padding[2] - dilation[2] * (K_w - 1) - 1) // stride[2] + 1 + mm_output_shape = (batch_size, D_out, H_out, W_out, C_out) + kernel_size = (K_d, K_h, K_w) + + if padding_mode != "zeros": + input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode) + padding = (0,) * (conv_type if conv_type != 1 else 2) + elif conv_type == 3: + input = torch.nn.functional.pad(input, reversed_padding_repeated_twice) + + if conv_type == 1: + input = input.unsqueeze(2) + + if conv_type == 3: + K_D_eff = K_d + (K_d - 1) * (dilation[0] - 1) + K_H_eff = K_h + (K_h - 1) * (dilation[0] - 1) + K_W_eff = K_w + (K_w - 1) * (dilation[0] - 1) + input = input.unfold(2, K_D_eff, stride[0]).unfold(3, K_H_eff, stride[1]).unfold(4, K_W_eff, stride[2]) + if dilation[0] > 1: + input = input[..., ::dilation[0], :, :] + if dilation[1] > 1: + input = input[..., ::dilation[1], :] + if dilation[2] > 1: + input = input[..., ::dilation[2]] + input = input.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(batch_size, D_out * H_out * W_out, -1) + else: + input = torch.nn.functional.unfold(input, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation).transpose(1,2) + return input, mm_output_shape + + +def quantized_conv_forward(self, input) -> torch.FloatTensor: + return self._conv_forward(input, self.sdnq_dequantizer(self.weight), self.bias) + + +def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: + output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation) + return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) + + +def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: + output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation) + return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) + + +def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: + output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation) + return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) diff --git a/modules/sdnq/layers/linear/forward.py b/modules/sdnq/layers/linear/forward.py new file mode 100644 index 000000000..20224204c --- /dev/null +++ b/modules/sdnq/layers/linear/forward.py @@ -0,0 +1,7 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +import torch + + +def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor: + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias) diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py new file mode 100644 index 000000000..17cac32e5 --- /dev/null +++ b/modules/sdnq/layers/linear/linear_fp8.py @@ -0,0 +1,41 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import Tuple + +import torch + +from ...common import use_torch_compile # noqa: TID252 + + +def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]: + input = input.flatten(0,-2).contiguous() + input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448) + input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn) + input_scale = input_scale.to(dtype=torch.float32) + return input, input_scale + + +def fp8_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, +) -> torch.FloatTensor: + return_dtype = input.dtype + output_shape = list(input.shape) + output_shape[-1] = weight.shape[-1] + input, input_scale = quantize_fp8_matmul_input(input) + return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(output_shape) + + +def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: + if torch.numel(input) / input.shape[-1] < 32: + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + return fp8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale) + + +if use_torch_compile: + try: + fp8_matmul = torch.compile(fp8_matmul, fullgraph=True, dynamic=False) + except Exception: + pass diff --git a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py new file mode 100644 index 000000000..768719d18 --- /dev/null +++ b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py @@ -0,0 +1,48 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import Tuple + +import torch + +from ...common import use_torch_compile # noqa: TID252 +from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 + + +def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]: + input = input.flatten(0,-2).contiguous() + input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448) + input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn) + scale = torch.mul(input_scale, scale) + if scale.dtype == torch.float16: # fp16 will overflow + scale = scale.to(dtype=torch.float32) + return input, scale + + +def fp8_matmul_tensorwise( + input: torch.FloatTensor, + weight: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, +) -> torch.FloatTensor: + return_dtype = input.dtype + output_shape = list(input.shape) + output_shape[-1] = weight.shape[-1] + dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) + input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) + if bias is not None: + return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, bias, return_dtype, output_shape) + else: + return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape) + + +def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor: + if torch.numel(input) / input.shape[-1] < 32: + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_dequantizer.scale) + + +if use_torch_compile: + try: + fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True, dynamic=False) + except Exception: + pass diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py new file mode 100644 index 000000000..1be08a5ae --- /dev/null +++ b/modules/sdnq/layers/linear/linear_int8.py @@ -0,0 +1,52 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import Tuple + +import torch + +from ...common import use_torch_compile # noqa: TID252 +from ...packed_int import unpack_int_symetric # noqa: TID252 +from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 + + +def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]: + input = input.flatten(0,-2).contiguous() + input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(127) + input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(dtype=torch.int8) + scale = torch.mul(input_scale, scale) + if scale.dtype == torch.float16: # fp16 will overflow + scale = scale.to(dtype=torch.float32) + return input, scale + + +def int8_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, + quantized_weight_shape: torch.Size, + weights_dtype: str, +) -> torch.FloatTensor: + if quantized_weight_shape is not None: + weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True) + return_dtype = input.dtype + output_shape = list(input.shape) + output_shape[-1] = weight.shape[-1] + input, scale = quantize_int8_matmul_input(input, scale) + if bias is not None: + return dequantize_symmetric_with_bias(torch._int_mm(input, weight), scale, bias, return_dtype, output_shape) + else: + return dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape) + + +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_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + return int8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale, getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), self.sdnq_dequantizer.weights_dtype) + + +if use_torch_compile: + try: + int8_matmul = torch.compile(int8_matmul, fullgraph=True, dynamic=False) + except Exception: + pass diff --git a/modules/sdnq/packed_int.py b/modules/sdnq/packed_int.py index b20c61818..84931d159 100644 --- a/modules/sdnq/packed_int.py +++ b/modules/sdnq/packed_int.py @@ -1,6 +1,7 @@ # pylint: disable=redefined-builtin,no-member,protected-access from typing import Optional + import torch from .common import dtype_dict