diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index a5bfd8194..b265fcce2 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -138,7 +138,7 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G shared.log.error(f'Network load: type=LoRA quant=bnb cls={self.__class__.__name__} type={self.quant_type} blocksize={self.blocksize} state={vars(self.quant_state)} weight={self.weight} bias={lora_weights} {e}') elif not bias and hasattr(self, "sdnq_decompressor"): try: - from modules.model_quant_sdnq import sdnq_quantize_layer + from modules.sdnq import sdnq_quantize_layer if hasattr(self, "sdnq_decompressor_backup"): sdnq_decompressor = self.sdnq_decompressor_backup.to(devices.device) else: diff --git a/modules/model_quant.py b/modules/model_quant.py index 099c175bc..6e2429e35 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -107,7 +107,7 @@ def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Mo from modules import shared if len(shared.opts.sdnq_quantize_weights) > 0 and (shared.opts.sdnq_quantize_mode == 'pre') and allow_sdnq: if 'Model' in shared.opts.sdnq_quantize_weights or (module is not None and module in shared.opts.sdnq_quantize_weights) or module == 'any': - from modules.model_quant_sdnq import SDNQQuantizer, SDNQConfig + from modules.sdnq import SDNQQuantizer, SDNQConfig diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig @@ -303,13 +303,13 @@ def apply_layerwise(sd_model, quiet:bool=False): def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True): global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement from modules import devices, shared - from modules.model_quant_sdnq import apply_sdnq_to_module + from modules.sdnq import apply_sdnq_to_module model.eval() if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: import torch - from modules.model_quant_sdnq import SDNQ_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 + from modules.sdnq import SDNQ_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 for i in range(len(model.encoder.block)): model.encoder.block[i].layer[1].DenseReluDense = SDNQ_T5DenseGatedActDense( model.encoder.block[i].layer[1].DenseReluDense, diff --git a/modules/model_quant_sdnq.py b/modules/model_quant_sdnq.py deleted file mode 100644 index 06e50ed29..000000000 --- a/modules/model_quant_sdnq.py +++ /dev/null @@ -1,1230 +0,0 @@ -# pylint: disable=redefined-builtin,no-member,protected-access - -from typing import Any, Dict, List, Tuple, Optional, Union -from dataclasses import dataclass -from enum import Enum -import sys -import torch -from diffusers.quantizers.base import DiffusersQuantizer -from diffusers.quantizers.quantization_config import QuantizationConfigMixin -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}, - "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}, - "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}, - "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}, - "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}, - "uint1": {"min": 0, "max": 1, "num_bits": 1, "target_dtype": torch.bool, "torch_dtype": torch.bool, "storage_dtype": torch.bool, "is_unsigned": True, "is_integer": True}, - "float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False}, - "float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False}, - "float8_e4m3fnuz": {"min": -240, "max": 240, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False}, - "float8_e5m2fnuz": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False}, -} -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") -if devices.backend in {"cpu", "openvino"}: - quantized_matmul_dtypes += ("float8_e4m3fnuz", "float8_e5m2fnuz") - -linear_types = ("Linear",) -conv_types = ("Conv1d", "Conv2d", "Conv3d") -conv_transpose_types = ("ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d") -allowed_types = linear_types + conv_types + conv_transpose_types - - -class QuantizationMethod(str, Enum): - SDNQ = "sdnq" - - -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, param_name=None, pre_mode=False): - layer_class_name = layer.__class__.__name__ - if layer_class_name in allowed_types: - is_conv_type = False - is_conv_transpose_type = False - is_linear_type = False - result_shape = None - if torch_dtype is None: - torch_dtype = devices.dtype - - if layer_class_name in conv_types: - if not quant_conv: - return layer - if dtype_dict[weights_dtype]["num_bits"] < 4: - weights_dtype = "uint4" - is_conv_type = True - reduction_axes = 1 - output_channel_size, channel_size = layer.weight.shape[:2] - group_channel_size = channel_size // layer.groups - use_quantized_matmul = False - if use_quantized_matmul_conv: - use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and group_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 group_channel_size % 16 == 0 - if use_quantized_matmul: - result_shape = layer.weight.shape - layer.weight.data = layer.weight.reshape(output_channel_size, -1) - elif layer_class_name in conv_transpose_types: - if not quant_conv: - return layer - if dtype_dict[weights_dtype]["num_bits"] < 4: - weights_dtype = "uint4" - is_conv_transpose_type = True - reduction_axes = 0 - channel_size, output_channel_size = layer.weight.shape[:2] - use_quantized_matmul = False - else: - is_linear_type = True - reduction_axes = -1 - output_channel_size, channel_size = layer.weight.shape - if use_quantized_matmul: - 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 - - if group_size == 0: - if is_linear_type: - group_size = 2 ** (2 + dtype_dict[weights_dtype]["num_bits"]) - else: - group_size = 2 ** (1 + dtype_dict[weights_dtype]["num_bits"]) - - 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 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] = (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"}: - return_device = devices.device - elif pre_mode: - if shared.opts.device_map == "gpu": - return_device = devices.device - elif shared.opts.sdnq_quantize_with_gpu: - return_device = devices.cpu - else: - return_device = layer.weight.device - else: - return_device = layer.weight.device - if not pre_mode: - if shared.opts.sdnq_quantize_with_gpu: - layer.weight.data = layer.weight.to(devices.device).to(dtype=torch.float32) - else: - layer.weight.data = layer.weight.to(dtype=torch.float32) - - if dtype_dict[weights_dtype]["is_unsigned"]: - scale, zero_point = get_scale_asymmetric(layer.weight, reduction_axes, weights_dtype) - else: - scale = get_scale_symmetric(layer.weight, reduction_axes, weights_dtype) - 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"] and not use_tensorwise_fp8_matmul): - scale = scale.to(torch_dtype) - if zero_point is not None: - zero_point = zero_point.to(torch_dtype) - - if use_quantized_matmul: - scale = scale.transpose(0,1) - if dtype_dict[weights_dtype]["num_bits"] == 8: - layer.weight.data = layer.weight.transpose(0,1) - if not dtype_dict[weights_dtype]["is_integer"]: - stride = layer.weight.stride() - if stride[0] > stride[1] and stride[1] == 1: - layer.weight.data = layer.weight.t().contiguous().t() - if not use_tensorwise_fp8_matmul: - scale = scale.to(torch.float32) - - layer.sdnq_decompressor = decompressor_dict[weights_dtype]( - scale=scale, - zero_point=zero_point, - compressed_weight_shape=layer.weight.shape, - result_dtype=torch_dtype, - result_shape=result_shape, - weights_dtype=weights_dtype, - use_quantized_matmul=use_quantized_matmul, - ) - layer.weight.data = layer.sdnq_decompressor.pack_weight(layer.weight).to(return_device) - layer.sdnq_decompressor = layer.sdnq_decompressor.to(return_device) - - if is_linear_type: - if use_quantized_matmul: - if dtype_dict[weights_dtype]["is_integer"]: - layer.forward = quantized_linear_forward_int8_matmul - else: - if use_tensorwise_fp8_matmul: - layer.forward = quantized_linear_forward_fp8_matmul_tensorwise - else: - layer.forward = quantized_linear_forward_fp8_matmul - else: - layer.forward = quantized_linear_forward - elif is_conv_type: - if use_quantized_matmul: - if dtype_dict[weights_dtype]["is_integer"]: - layer.forward = quantized_conv_forward_int8_matmul - else: - if use_tensorwise_fp8_matmul: - layer.forward = quantized_conv_forward_fp8_matmul_tensorwise - else: - layer.forward = quantized_conv_forward_fp8_matmul - else: - layer.forward = quantized_conv_forward - elif is_conv_transpose_type: - if layer_class_name.endswith("1d"): - layer.forward = quantized_conv_transpose_1d_forward - elif layer_class_name.endswith("2d"): - layer.forward = quantized_conv_transpose_2d_forward - elif layer_class_name.endswith("3d"): - layer.forward = quantized_conv_transpose_3d_forward - layer.forward = layer.forward.__get__(layer, layer.__class__) - devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}") - return layer - - -def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, param_name=None): - has_children = list(model.children()) - if not has_children: - return model - for module_param_name, module in model.named_children(): - if hasattr(module, "weight") and module.weight is not None: - module = sdnq_quantize_layer( - module, - weights_dtype=weights_dtype, - torch_dtype=torch_dtype, - group_size=group_size, - quant_conv=quant_conv, - use_quantized_matmul=use_quantized_matmul, - use_quantized_matmul_conv=use_quantized_matmul_conv, - param_name=module_param_name, - ) - module = apply_sdnq_to_module( - module, - weights_dtype=weights_dtype, - torch_dtype=torch_dtype, - group_size=group_size, - quant_conv=quant_conv, - use_quantized_matmul=use_quantized_matmul, - use_quantized_matmul_conv=use_quantized_matmul_conv, - param_name=module_param_name, - ) - return model - - -def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: 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: List[int], weights_dtype: str) -> torch.FloatTensor: - abs_min_values = torch.amin(weight, dim=reduction_axes, keepdims=True).abs_() - max_values = torch.amax(weight, dim=reduction_axes, keepdims=True) - scale = torch.where(abs_min_values >= max_values, abs_min_values, -max_values).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, scale: torch.FloatTensor, zero_point: torch.FloatTensor, weights_dtype: str) -> torch.ByteTensor: - if zero_point is not None: - compressed_weight = torch.sub(weight, zero_point).div_(scale) - else: - compressed_weight = torch.div(weight, scale) - if dtype_dict[weights_dtype]["is_integer"]: - compressed_weight.round_() - compressed_weight = compressed_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"]) - return compressed_weight - - -def decompress_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.Tensor: - result = torch.addcmul(zero_point, input.to(dtype=scale.dtype), scale).to(dtype=dtype) - if result_shape is not None: - result = result.reshape(result_shape) - return result - - -def decompress_symmetric(input: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.Tensor: - if skip_quantized_matmul: - result = input.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype) - else: - result = input.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype) - if result_shape is not None: - result = result.reshape(result_shape) - return result - - -def decompress_packed_int_asymmetric(input: torch.Tensor, scale: torch.Tensor, zero_point: torch.Tensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.Tensor: - return decompress_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape), scale, zero_point, dtype, result_shape) - - -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_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape) - else: - return decompress_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape) - - -def pack_int_symetric(tensor: torch.Tensor, weights_dtype: str) -> torch.Tensor: - return packed_int_function_dict[weights_dtype]["pack"](tensor.to(dtype=dtype_dict[weights_dtype]["torch_dtype"]).sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"])) - - -def unpack_int_symetric(packed_tensor: torch.Tensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, transpose: Optional[bool] = False) -> torch.Tensor: - if dtype is None: - dtype = dtype_dict[weights_dtype]["torch_dtype"] - result = packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"]) - if transpose: - result = result.transpose(0,1) - return result - - -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(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 2), 192)), - torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 4), 192)), - torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 6)), - ), - dim=-1 - ) - return packed_tensor - - -def pack_uint5(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_left_shift(packed_tensor[:, 5], 5)), - torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_left_shift(packed_tensor[:, 6], 5)), - torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 7], 5)), - torch.bitwise_or( - packed_tensor[:, 3], - torch.bitwise_or( - torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 5], 2), 96), - torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128), - ), - ), - torch.bitwise_or( - packed_tensor[:, 4], - torch.bitwise_or( - torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 2), 96), - torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128), - ), - ), - ), - dim=-1 - ) - return packed_tensor - - -def unpack_uint5(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: - result = torch.stack( - ( - torch.bitwise_and(packed_tensor[:, 0], 31), - torch.bitwise_and(packed_tensor[:, 1], 31), - torch.bitwise_and(packed_tensor[:, 2], 31), - torch.bitwise_and(packed_tensor[:, 3], 31), - torch.bitwise_and(packed_tensor[:, 4], 31), - torch.bitwise_or( - torch.bitwise_right_shift(packed_tensor[:, 0], 5), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 2), 24), - ), - torch.bitwise_or( - torch.bitwise_right_shift(packed_tensor[:, 1], 5), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 2), 24), - ), - torch.bitwise_or( - torch.bitwise_right_shift(packed_tensor[:, 2], 5), - torch.bitwise_or( - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 3), 16), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 4), 8), - ), - ), - ), - dim=-1 - ).reshape(shape) - return result - - - -def pack_uint4(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, 2) - packed_tensor = torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 4)) - return packed_tensor - - -def pack_uint3(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( - torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 3)), - torch.bitwise_left_shift(packed_tensor[:, 6], 6), - ), - torch.bitwise_or( - torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 3)), - torch.bitwise_left_shift(packed_tensor[:, 7], 6), - ), - torch.bitwise_or( - torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_left_shift(packed_tensor[:, 5], 3)), - torch.bitwise_or( - torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 4), 64), - torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128), - ) - ), - ), - dim=-1 - ) - return packed_tensor - - -def pack_uint2(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.bitwise_or( - torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 2)), - torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 2], 4), torch.bitwise_left_shift(packed_tensor[:, 3], 6)), - ) - 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_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 - - -def unpack_uint3(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: - result = torch.stack( - ( - torch.bitwise_and(packed_tensor[:, 0], 7), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 3), 7), - torch.bitwise_and(packed_tensor[:, 1], 7), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 3), 7), - torch.bitwise_and(packed_tensor[:, 2], 7), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 7), - torch.bitwise_or( - torch.bitwise_right_shift(packed_tensor[:, 0], 6), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 4), 4), - ), - torch.bitwise_or( - torch.bitwise_right_shift(packed_tensor[:, 1], 6), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 5), 4), - ), - ), - dim=-1 - ).reshape(shape) - return result - - -def unpack_uint2(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: - result = torch.stack( - ( - torch.bitwise_and(packed_tensor, 3), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 2), 3), - torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 4), 3), - torch.bitwise_right_shift(packed_tensor, 6), - ), - dim=-1 - ).reshape(shape) - return result - - -def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.FloatTensor, torch.FloatTensor]: - input = input.flatten(0,-2).contiguous() - input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 448) - input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn) - input_scale = input_scale.to(torch.float32) - return input, input_scale - - -def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.ByteTensor, torch.FloatTensor]: - input = input.flatten(0,-2).contiguous() - input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 448) - input = torch.div(input, input_scale).clamp_(-448, 448).to(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.ByteTensor, torch.FloatTensor]: - input = input.flatten(0,-2).contiguous() - input_scale = torch.div(input.abs().amax(dim=-1, keepdims=True), 127) - input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(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) - result = decompress_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) - if bias is not None: - result.add_(bias) - return result - - -def int8_matmul( - input: torch.FloatTensor, - weight: torch.Tensor, - 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_int_symetric(weight, compressed_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) - result = decompress_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape) - if bias is not None: - result.add_(bias) - return result - - -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) - elif conv_type == 3: - 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.ByteTensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - result_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, 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 = [] - 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 bias is not None: - result.add_(bias) - - 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.ByteTensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - result_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_fp8_matmul_input_tensorwise(input, scale) - dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) - - if groups == 1: - result = decompress_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, mm_output_shape) - 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 = decompress_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape) - if bias is not None: - result.add_(bias) - - 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.ByteTensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - result_shape: torch.Size, - compressed_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 compressed_weight_shape is not None: - weight = unpack_int_symetric(weight, compressed_weight_shape, weights_dtype, dtype=torch.int8, transpose=True) - - if groups == 1: - result = decompress_symmetric(torch._int_mm(input, weight), scale, return_dtype, mm_output_shape) - 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 = decompress_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape) - if bias is not None: - result.add_(bias) - - 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_decompressor(self.weight, skip_quantized_matmul=True), self.bias) - return fp8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.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_decompressor(self.weight, skip_quantized_matmul=True), self.bias) - return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_decompressor.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_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), self.sdnq_decompressor.weights_dtype) - - -def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor: - return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight), self.bias) - - -def get_conv_args(input_ndim, stride, padding, dilation): - if input_ndim == 3: - conv_type = 1 - elif input_ndim == 4: - conv_type = 2 - elif input_ndim == 5: - 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_decompressor(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_decompressor.scale, - self.sdnq_decompressor.result_shape, - self.sdnq_decompressor.weights_dtype, - 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_decompressor(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_decompressor.scale, - self.sdnq_decompressor.result_shape, - self.sdnq_decompressor.weights_dtype, - 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_decompressor(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_decompressor.scale, - self.sdnq_decompressor.result_shape, - getattr(self.sdnq_decompressor, "compressed_weight_shape", None), - self.sdnq_decompressor.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_decompressor(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_decompressor(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_decompressor(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_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) - - -class AsymmetricWeightsDecompressor(torch.nn.Module): - def __init__( - self, - scale: torch.Tensor, - zero_point: torch.Tensor, - result_dtype: torch.dtype, - result_shape: torch.Size, - weights_dtype: str, - **kwargs, - ): - super().__init__() - self.weights_dtype = weights_dtype - self.use_quantized_matmul = False - self.result_dtype = result_dtype - self.result_shape = result_shape - self.register_buffer("scale", scale) - self.register_buffer("zero_point", zero_point) - - def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: - return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]) - - def forward(self, weight, **kwargs): - return decompress_asymmetric_compiled(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape) - - -class SymmetricWeightsDecompressor(torch.nn.Module): - def __init__( - self, - scale: torch.Tensor, - result_dtype: torch.dtype, - result_shape: torch.Size, - weights_dtype: str, - use_quantized_matmul: bool = False, - **kwargs, - ): - super().__init__() - self.weights_dtype = weights_dtype - self.use_quantized_matmul = use_quantized_matmul - self.result_dtype = result_dtype - self.result_shape = result_shape - self.register_buffer("scale", scale) - - def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: - return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]) - - def forward(self, weight, skip_quantized_matmul=False, **kwargs): - return decompress_symmetric_compiled(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul) - - -class PackedINTAsymmetricWeightsDecompressor(torch.nn.Module): - def __init__( - self, - scale: torch.Tensor, - zero_point: torch.Tensor, - compressed_weight_shape: torch.Size, - result_dtype: torch.dtype, - result_shape: torch.Size, - weights_dtype: str, - **kwargs, - ): - super().__init__() - self.weights_dtype = weights_dtype - self.use_quantized_matmul = False - self.compressed_weight_shape = compressed_weight_shape - self.result_dtype = result_dtype - self.result_shape = result_shape - self.register_buffer("scale", scale) - self.register_buffer("zero_point", zero_point) - - def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: - return packed_int_function_dict[self.weights_dtype]["pack"](weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])) - - def forward(self, weight, **kwargs): - 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 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 = weights_dtype - self.use_quantized_matmul = use_quantized_matmul - self.compressed_weight_shape = compressed_weight_shape - self.result_dtype = result_dtype - self.result_shape = result_shape - self.register_buffer("scale", scale) - - def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: - return pack_int_symetric(weight, self.weights_dtype) - - def forward(self, weight, skip_quantized_matmul=False, **kwargs): - 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, - "int6": PackedINTSymmetricWeightsDecompressor, - "int5": PackedINTSymmetricWeightsDecompressor, - "int4": PackedINTSymmetricWeightsDecompressor, - "int3": PackedINTSymmetricWeightsDecompressor, - "int2": PackedINTSymmetricWeightsDecompressor, - "uint8": AsymmetricWeightsDecompressor, - "uint6": PackedINTAsymmetricWeightsDecompressor, - "uint5": PackedINTAsymmetricWeightsDecompressor, - "uint4": PackedINTAsymmetricWeightsDecompressor, - "uint3": PackedINTAsymmetricWeightsDecompressor, - "uint2": PackedINTAsymmetricWeightsDecompressor, - "uint1": AsymmetricWeightsDecompressor, - "bool": AsymmetricWeightsDecompressor, - "float8_e4m3fn": SymmetricWeightsDecompressor, - "float8_e4m3fnuz": SymmetricWeightsDecompressor, - "float8_e5m2": SymmetricWeightsDecompressor, - "float8_e5m2fnuz": SymmetricWeightsDecompressor, -} - - -packed_int_function_dict = { - "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}, - "uint6": {"pack": pack_uint6, "unpack": unpack_uint6}, - "uint5": {"pack": pack_uint5, "unpack": unpack_uint5}, - "uint4": {"pack": pack_uint4, "unpack": unpack_uint4}, - "uint3": {"pack": pack_uint3, "unpack": unpack_uint3}, - "uint2": {"pack": pack_uint2, "unpack": unpack_uint2}, -} - - -class SDNQQuantizer(DiffusersQuantizer): - r""" - Diffusers Quantizer for SDNQ - """ - - requires_parameters_quantization = True - use_keep_in_fp32_modules = True - requires_calibration = False - required_packages = None - torch_dtype = None - - def __init__(self, quantization_config, **kwargs): # pylint: disable=useless-parent-delegation - super().__init__(quantization_config, **kwargs) - - def check_if_quantized_param( - self, - model, - param_value: "torch.Tensor", - param_name: str, - state_dict: Dict[str, Any], - **kwargs, - ): - if param_name.endswith(".weight"): - split_param_name = param_name.split(".") - if param_name not in self.modules_to_not_convert and not any(param in split_param_name for param in self.modules_to_not_convert): - layer_class_name = get_module_from_name(model, param_name)[0].__class__.__name__ - if layer_class_name in allowed_types: - if layer_class_name in conv_types or layer_class_name in conv_transpose_types: - if self.quantization_config.quant_conv: - return True - else: - return True - param_value.data = param_value.clone() # safetensors is unable to release the cpu memory without this - return False - - def check_quantized_param(self, *args, **kwargs) -> bool: - """ - needed for transformers compatibilty, returns self.check_if_quantized_param - """ - return self.check_if_quantized_param(*args, **kwargs) - - def create_quantized_param( # pylint: disable=arguments-differ - self, - model, - param_value: torch.FloatTensor, - param_name: str, - target_device: torch.device, - state_dict: Dict[str, Any], # pylint: disable=unused-argument - unexpected_keys: List[str], # pylint: disable=unused-argument - **kwargs, - ): - # load the model params to target_device first - layer, _ = get_module_from_name(model, param_name) - if shared.opts.sdnq_quantize_with_gpu: - if param_value.dtype == torch.float32 and devices.same_device(param_value.device, devices.device): - param_value = param_value.clone() - else: - param_value = param_value.to(devices.device).to(dtype=torch.float32) - else: - if param_value.dtype == torch.float32 and devices.same_device(param_value.device, target_device): - param_value = param_value.clone() - else: - param_value = param_value.to(target_device).to(dtype=torch.float32) - layer.weight = torch.nn.Parameter(param_value, requires_grad=False) - layer = sdnq_quantize_layer( - layer, - weights_dtype=self.quantization_config.weights_dtype, - torch_dtype=self.torch_dtype, - group_size=self.quantization_config.group_size, - quant_conv=self.quantization_config.quant_conv, - use_quantized_matmul=self.quantization_config.use_quantized_matmul, - use_quantized_matmul_conv=self.quantization_config.use_quantized_matmul_conv, - param_name=param_name, - pre_mode=True, - ) - - def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: - max_memory = {key: val * 0.80 for key, val in max_memory.items()} - return max_memory - - def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed - return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"] - - def update_torch_dtype(self, torch_dtype: torch.dtype = None) -> torch.dtype: - if torch_dtype is None: - torch_dtype = devices.dtype - self.torch_dtype = torch_dtype - return torch_dtype - - def _process_model_before_weight_loading( # pylint: disable=arguments-differ - self, - model, - device_map, # pylint: disable=unused-argument - keep_in_fp32_modules: List[str] = [], - **kwargs, - ): - model.config.quantization_config = self.quantization_config - self.modules_to_not_convert = self.quantization_config.modules_to_not_convert - if not isinstance(self.modules_to_not_convert, list): - self.modules_to_not_convert = [self.modules_to_not_convert] - if keep_in_fp32_modules is not None: - self.modules_to_not_convert.extend(keep_in_fp32_modules) - - def _process_model_after_weight_loading(self, model, **kwargs): - if shared.opts.diffusers_offload_mode != "none": - model = model.to(devices.cpu) - devices.torch_gc(force=True) - return model - - def get_cuda_warm_up_factor(self): - return 32 // dtype_dict[self.quantization_config.weights_dtype]["num_bits"] - - def update_tp_plan(self, config): - """ - needed for transformers compatibilty, no-op function - """ - return config - - def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument - """ - needed for transformers compatibilty, no-op function - """ - return unexpected_keys - - def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument - """ - needed for transformers compatibilty, no-op function - """ - return missing_keys - - def update_expected_keys(self, model, expected_keys: List[str], loaded_keys: List[str]) -> List[str]: # pylint: disable=unused-argument - """ - needed for transformers compatibilty, no-op function - """ - return expected_keys - - @property - def is_trainable(self): - return False - - @property - def is_serializable(self): - return False - - -@dataclass -class SDNQConfig(QuantizationConfigMixin): - """ - This is a wrapper class about all possible attributes and features that you can play with a model that has been - loaded using `sdnq`. - - Args: - weights_dtype (`str`, *optional*, defaults to `"int8"`): - The target dtype for the weights after quantization. Supported values are: - ("int8", "int6", "int5", "int4", "int3", "int2", "uint8", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz") - modules_to_not_convert (`list`, *optional*, default to `None`): - The list of modules to not quantize, useful for quantizing models that explicitly require to have some - modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers). - """ - - def __init__( # pylint: disable=super-init-not-called - self, - weights_dtype: str = "int8", - group_size: int = 0, - quant_conv: bool = False, - use_quantized_matmul: bool = False, - use_quantized_matmul_conv: bool = False, - modules_to_not_convert: Optional[List[str]] = None, - **kwargs, # pylint: disable=unused-argument - ): - self.weights_dtype = weights_dtype - self.quant_method = QuantizationMethod.SDNQ - self.group_size = group_size - self.quant_conv = quant_conv - self.use_quantized_matmul = use_quantized_matmul - self.use_quantized_matmul_conv = use_quantized_matmul_conv - self.modules_to_not_convert = modules_to_not_convert - self.post_init() - self.is_integer = dtype_dict[self.weights_dtype]["is_integer"] - - def post_init(self): - r""" - Safety checker that arguments are correct - """ - accepted_weights = ["int8", "int6", "int5", "int4", "int3", "int2", "uint8", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz"] - if self.weights_dtype not in accepted_weights: - raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}") - - -class SDNQ_T5DenseGatedActDense(torch.nn.Module): # forward can't find what self is without creating a class - def __init__(self, T5DenseGatedActDense, dtype): - super().__init__() - self.wi_0 = T5DenseGatedActDense.wi_0 - self.wi_1 = T5DenseGatedActDense.wi_1 - self.wo = T5DenseGatedActDense.wo - self.dropout = T5DenseGatedActDense.dropout - self.act = T5DenseGatedActDense.act - self.torch_dtype = dtype - - def forward(self, hidden_states): - hidden_gelu = self.act(self.wi_0(hidden_states)) - hidden_linear = self.wi_1(hidden_states) - hidden_states = hidden_gelu * hidden_linear - hidden_states = self.dropout(hidden_states) - hidden_states = hidden_states.to(self.torch_dtype) # this line needs to be forced - hidden_states = self.wo(hidden_states) - return hidden_states - - -if shared.opts.sdnq_decompress_compile: - try: - torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) - 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_packed_int_symmetric_compiled = torch.compile(decompress_packed_int_symmetric, fullgraph=True) - int8_matmul = torch.compile(int8_matmul, fullgraph=True) - fp8_matmul = torch.compile(fp8_matmul, fullgraph=True) - fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True) - conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True) - conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True) - conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, 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_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_packed_int_symmetric_compiled = decompress_packed_int_symmetric diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py new file mode 100644 index 000000000..fde0b2a26 --- /dev/null +++ b/modules/sdnq/__init__.py @@ -0,0 +1,428 @@ +# pylint: disable=redefined-builtin,no-member,protected-access + +from typing import Any, Dict, List, Tuple, Optional, Union +from dataclasses import dataclass +from enum import Enum +import torch +from diffusers.quantizers.base import DiffusersQuantizer +from diffusers.quantizers.quantization_config import QuantizationConfigMixin +from diffusers.utils import get_module_from_name +from modules import devices, shared + +from .common import dtype_dict, use_tensorwise_fp8_matmul, quantized_matmul_dtypes, allowed_types, conv_types, conv_transpose_types +from .decompressor import decompressor_dict +from .forward import get_forward_func + + +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, param_name=None, pre_mode=False): + layer_class_name = layer.__class__.__name__ + if layer_class_name in allowed_types: + is_conv_type = False + is_conv_transpose_type = False + is_linear_type = False + result_shape = None + if torch_dtype is None: + torch_dtype = devices.dtype + + if layer_class_name in conv_types: + if not quant_conv: + return layer + if dtype_dict[weights_dtype]["num_bits"] < 4: + weights_dtype = "uint4" + is_conv_type = True + reduction_axes = 1 + output_channel_size, channel_size = layer.weight.shape[:2] + group_channel_size = channel_size // layer.groups + use_quantized_matmul = False + if use_quantized_matmul_conv: + use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and group_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 group_channel_size % 16 == 0 + if use_quantized_matmul: + result_shape = layer.weight.shape + layer.weight.data = layer.weight.reshape(output_channel_size, -1) + elif layer_class_name in conv_transpose_types: + if not quant_conv: + return layer + if dtype_dict[weights_dtype]["num_bits"] < 4: + weights_dtype = "uint4" + is_conv_transpose_type = True + reduction_axes = 0 + channel_size, output_channel_size = layer.weight.shape[:2] + use_quantized_matmul = False + else: + is_linear_type = True + reduction_axes = -1 + output_channel_size, channel_size = layer.weight.shape + if use_quantized_matmul: + 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 + + if group_size == 0: + if is_linear_type: + group_size = 2 ** (2 + dtype_dict[weights_dtype]["num_bits"]) + else: + group_size = 2 ** (1 + dtype_dict[weights_dtype]["num_bits"]) + + 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 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] = (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"}: + return_device = devices.device + elif pre_mode: + if shared.opts.device_map == "gpu": + return_device = devices.device + elif shared.opts.sdnq_quantize_with_gpu: + return_device = devices.cpu + else: + return_device = layer.weight.device + else: + return_device = layer.weight.device + if not pre_mode: + if shared.opts.sdnq_quantize_with_gpu: + layer.weight.data = layer.weight.to(devices.device).to(dtype=torch.float32) + else: + layer.weight.data = layer.weight.to(dtype=torch.float32) + + if dtype_dict[weights_dtype]["is_unsigned"]: + scale, zero_point = get_scale_asymmetric(layer.weight, reduction_axes, weights_dtype) + layer.weight.data.sub_(zero_point).div_(scale) + else: + scale = get_scale_symmetric(layer.weight, reduction_axes, weights_dtype) + layer.weight.data.div_(scale) + zero_point = None + if dtype_dict[weights_dtype]["is_integer"]: + layer.weight.data.round_() + layer.weight.data = layer.weight.data.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"]) + + 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) + + if use_quantized_matmul: + scale = scale.transpose(0,1) + if dtype_dict[weights_dtype]["num_bits"] == 8: + layer.weight.data = layer.weight.transpose(0,1) + if not dtype_dict[weights_dtype]["is_integer"]: + stride = layer.weight.stride() + if stride[0] > stride[1] and stride[1] == 1: + layer.weight.data = layer.weight.t().contiguous().t() + if not use_tensorwise_fp8_matmul: + scale = scale.to(torch.float32) + + layer.sdnq_decompressor = decompressor_dict[weights_dtype]( + scale=scale, + zero_point=zero_point, + compressed_weight_shape=layer.weight.shape, + result_dtype=torch_dtype, + result_shape=result_shape, + weights_dtype=weights_dtype, + use_quantized_matmul=use_quantized_matmul, + ) + layer.weight.data = layer.sdnq_decompressor.pack_weight(layer.weight).to(return_device) + layer.sdnq_decompressor = layer.sdnq_decompressor.to(return_device) + + 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 + + +def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, param_name=None): # pylint: disable=unused-argument + has_children = list(model.children()) + if not has_children: + return model + for module_param_name, module in model.named_children(): + if hasattr(module, "weight") and module.weight is not None: + module = sdnq_quantize_layer( + module, + weights_dtype=weights_dtype, + torch_dtype=torch_dtype, + group_size=group_size, + quant_conv=quant_conv, + use_quantized_matmul=use_quantized_matmul, + use_quantized_matmul_conv=use_quantized_matmul_conv, + param_name=module_param_name, + ) + module = apply_sdnq_to_module( + module, + weights_dtype=weights_dtype, + torch_dtype=torch_dtype, + group_size=group_size, + quant_conv=quant_conv, + use_quantized_matmul=use_quantized_matmul, + use_quantized_matmul_conv=use_quantized_matmul_conv, + param_name=module_param_name, + ) + return model + + +def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: 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: List[int], weights_dtype: str) -> torch.FloatTensor: + abs_min_values = torch.amin(weight, dim=reduction_axes, keepdims=True).abs_() + max_values = torch.amax(weight, dim=reduction_axes, keepdims=True) + scale = torch.where(abs_min_values >= max_values, abs_min_values, -max_values).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 + + +class QuantizationMethod(str, Enum): + SDNQ = "sdnq" + + +class SDNQQuantizer(DiffusersQuantizer): + r""" + Diffusers Quantizer for SDNQ + """ + + requires_parameters_quantization = True + use_keep_in_fp32_modules = True + requires_calibration = False + required_packages = None + torch_dtype = None + + def __init__(self, quantization_config, **kwargs): # pylint: disable=useless-parent-delegation + super().__init__(quantization_config, **kwargs) + self.modules_to_not_convert = [] + + def check_if_quantized_param( + self, + model, + param_value: "torch.Tensor", + param_name: str, + state_dict: Dict[str, Any], # pylint: disable=unused-argument + **kwargs, # pylint: disable=unused-argument + ): + if param_name.endswith(".weight"): + split_param_name = param_name.split(".") + if param_name not in self.modules_to_not_convert and not any(param in split_param_name for param in self.modules_to_not_convert): + layer_class_name = get_module_from_name(model, param_name)[0].__class__.__name__ + if layer_class_name in allowed_types: + if layer_class_name in conv_types or layer_class_name in conv_transpose_types: + if self.quantization_config.quant_conv: + return True + else: + return True + param_value.data = param_value.clone() # safetensors is unable to release the cpu memory without this + return False + + def check_quantized_param(self, *args, **kwargs) -> bool: + """ + needed for transformers compatibilty, returns self.check_if_quantized_param + """ + return self.check_if_quantized_param(*args, **kwargs) + + def create_quantized_param( # pylint: disable=arguments-differ + self, + model, + param_value: torch.FloatTensor, + param_name: str, + target_device: torch.device, + state_dict: Dict[str, Any], # pylint: disable=unused-argument + unexpected_keys: List[str], # pylint: disable=unused-argument + **kwargs, # pylint: disable=unused-argument + ): + # load the model params to target_device first + layer, _ = get_module_from_name(model, param_name) + if shared.opts.sdnq_quantize_with_gpu: + if param_value.dtype == torch.float32 and devices.same_device(param_value.device, devices.device): + param_value = param_value.clone() + else: + param_value = param_value.to(devices.device).to(dtype=torch.float32) + else: + if param_value.dtype == torch.float32 and devices.same_device(param_value.device, target_device): + param_value = param_value.clone() + else: + param_value = param_value.to(target_device).to(dtype=torch.float32) + layer.weight = torch.nn.Parameter(param_value, requires_grad=False) + layer = sdnq_quantize_layer( + layer, + weights_dtype=self.quantization_config.weights_dtype, + torch_dtype=self.torch_dtype, + group_size=self.quantization_config.group_size, + quant_conv=self.quantization_config.quant_conv, + use_quantized_matmul=self.quantization_config.use_quantized_matmul, + use_quantized_matmul_conv=self.quantization_config.use_quantized_matmul_conv, + param_name=param_name, + pre_mode=True, + ) + + def adjust_max_memory(self, max_memory: Dict[str, Union[int, str]]) -> Dict[str, Union[int, str]]: + max_memory = {key: val * 0.80 for key, val in max_memory.items()} + return max_memory + + def adjust_target_dtype(self, target_dtype: torch.dtype) -> torch.dtype: # pylint: disable=unused-argument,arguments-renamed + return dtype_dict[self.quantization_config.weights_dtype]["target_dtype"] + + def update_torch_dtype(self, torch_dtype: torch.dtype = None) -> torch.dtype: + if torch_dtype is None: + torch_dtype = devices.dtype + self.torch_dtype = torch_dtype + return torch_dtype + + def _process_model_before_weight_loading( # pylint: disable=arguments-differ + self, + model, + device_map, # pylint: disable=unused-argument + keep_in_fp32_modules: List[str] = [], + **kwargs, # pylint: disable=unused-argument + ): + model.config.quantization_config = self.quantization_config + self.modules_to_not_convert.extend(self.quantization_config.modules_to_not_convert) + if keep_in_fp32_modules is not None: + self.modules_to_not_convert.extend(keep_in_fp32_modules) + + def _process_model_after_weight_loading(self, model, **kwargs): # pylint: disable=unused-argument + if shared.opts.diffusers_offload_mode != "none": + model = model.to(devices.cpu) + devices.torch_gc(force=True) + return model + + def get_cuda_warm_up_factor(self): + return 32 // dtype_dict[self.quantization_config.weights_dtype]["num_bits"] + + def update_tp_plan(self, config): + """ + needed for transformers compatibilty, no-op function + """ + return config + + def update_unexpected_keys(self, model, unexpected_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument + """ + needed for transformers compatibilty, no-op function + """ + return unexpected_keys + + def update_missing_keys_after_loading(self, model, missing_keys: List[str], prefix: str) -> List[str]: # pylint: disable=unused-argument + """ + needed for transformers compatibilty, no-op function + """ + return missing_keys + + def update_expected_keys(self, model, expected_keys: List[str], loaded_keys: List[str]) -> List[str]: # pylint: disable=unused-argument + """ + needed for transformers compatibilty, no-op function + """ + return expected_keys + + @property + def is_trainable(self): + return False + + @property + def is_serializable(self): + return False + + +@dataclass +class SDNQConfig(QuantizationConfigMixin): + """ + This is a wrapper class about all possible attributes and features that you can play with a model that has been + loaded using `sdnq`. + + Args: + weights_dtype (`str`, *optional*, defaults to `"int8"`): + The target dtype for the weights after quantization. Supported values are: + ("int8", "int6", "int5", "int4", "int3", "int2", "uint8", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz") + modules_to_not_convert (`list`, *optional*, default to `None`): + The list of modules to not quantize, useful for quantizing models that explicitly require to have some + modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers). + """ + + def __init__( # pylint: disable=super-init-not-called + self, + weights_dtype: str = "int8", + group_size: int = 0, + quant_conv: bool = False, + use_quantized_matmul: bool = False, + use_quantized_matmul_conv: bool = False, + modules_to_not_convert: Optional[List[str]] = None, + **kwargs, # pylint: disable=unused-argument + ): + self.weights_dtype = weights_dtype + self.quant_method = QuantizationMethod.SDNQ + self.group_size = group_size + self.quant_conv = quant_conv + self.use_quantized_matmul = use_quantized_matmul + self.use_quantized_matmul_conv = use_quantized_matmul_conv + self.modules_to_not_convert = modules_to_not_convert + self.post_init() + self.is_integer = dtype_dict[self.weights_dtype]["is_integer"] + + def post_init(self): + r""" + Safety checker that arguments are correct + """ + accepted_weights = ["int8", "int6", "int5", "int4", "int3", "int2", "uint8", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz"] + if self.weights_dtype not in accepted_weights: + raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}") + if not isinstance(self.modules_to_not_convert, list): + self.modules_to_not_convert = [self.modules_to_not_convert] + + +class SDNQ_T5DenseGatedActDense(torch.nn.Module): # forward can't find what self is without creating a class + def __init__(self, T5DenseGatedActDense, dtype): + super().__init__() + self.wi_0 = T5DenseGatedActDense.wi_0 + self.wi_1 = T5DenseGatedActDense.wi_1 + self.wo = T5DenseGatedActDense.wo + self.dropout = T5DenseGatedActDense.dropout + self.act = T5DenseGatedActDense.act + self.torch_dtype = dtype + + def forward(self, hidden_states): + hidden_gelu = self.act(self.wi_0(hidden_states)) + hidden_linear = self.wi_1(hidden_states) + hidden_states = hidden_gelu * hidden_linear + hidden_states = self.dropout(hidden_states) + hidden_states = hidden_states.to(self.torch_dtype) # this line needs to be forced + hidden_states = self.wo(hidden_states) + return hidden_states diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py new file mode 100644 index 000000000..70addbe7a --- /dev/null +++ b/modules/sdnq/common.py @@ -0,0 +1,41 @@ +# pylint: disable=redefined-builtin,no-member,protected-access + +import sys +import torch +from accelerate.utils import CustomDtype +from modules import devices + +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}, + "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}, + "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}, + "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}, + "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}, + "uint1": {"min": 0, "max": 1, "num_bits": 1, "target_dtype": torch.bool, "torch_dtype": torch.bool, "storage_dtype": torch.bool, "is_unsigned": True, "is_integer": True}, + "float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False}, + "float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False}, + "float8_e4m3fnuz": {"min": -240, "max": 240, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False}, + "float8_e5m2fnuz": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False}, +} +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") +if devices.backend in {"cpu", "openvino"}: + quantized_matmul_dtypes += ("float8_e4m3fnuz", "float8_e5m2fnuz") + +linear_types = ("Linear",) +conv_types = ("Conv1d", "Conv2d", "Conv3d") +conv_transpose_types = ("ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d") +allowed_types = linear_types + conv_types + conv_transpose_types diff --git a/modules/sdnq/decompressor.py b/modules/sdnq/decompressor.py new file mode 100644 index 000000000..2c37f26b2 --- /dev/null +++ b/modules/sdnq/decompressor.py @@ -0,0 +1,170 @@ +# pylint: disable=redefined-builtin,no-member,protected-access + +import torch +from modules import shared + +from .common import dtype_dict +from .packed_int import pack_int_symetric, unpack_int_symetric, packed_int_function_dict + + +def decompress_asymmetric(input: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor: + result = torch.addcmul(zero_point, input.to(dtype=scale.dtype), scale).to(dtype=dtype) + if result_shape is not None: + result = result.reshape(result_shape) + return result + + +def decompress_symmetric(input: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor: + if skip_quantized_matmul: + result = input.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype) + else: + result = input.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype) + if result_shape is not None: + result = result.reshape(result_shape) + return result + + +def decompress_packed_int_asymmetric(input: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor: + return decompress_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](input, shape), scale, zero_point, dtype, result_shape) + + +def decompress_packed_int_symmetric(input: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.FloatTensor: + if skip_quantized_matmul: + return decompress_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape) + else: + return decompress_symmetric(unpack_int_symetric(input, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape) + + +class AsymmetricWeightsDecompressor(torch.nn.Module): + def __init__( + self, + scale: torch.Tensor, + zero_point: torch.Tensor, + result_dtype: torch.dtype, + result_shape: torch.Size, + weights_dtype: str, + **kwargs, # pylint: disable=unused-argument + ): + super().__init__() + self.weights_dtype = weights_dtype + self.use_quantized_matmul = False + self.result_dtype = result_dtype + self.result_shape = result_shape + self.register_buffer("scale", scale) + self.register_buffer("zero_point", zero_point) + + def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: + return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]) + + def forward(self, weight, **kwargs): # pylint: disable=unused-argument + return decompress_asymmetric(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape) + + +class SymmetricWeightsDecompressor(torch.nn.Module): + def __init__( + self, + scale: torch.Tensor, + result_dtype: torch.dtype, + result_shape: torch.Size, + weights_dtype: str, + use_quantized_matmul: bool = False, + **kwargs, # pylint: disable=unused-argument + ): + super().__init__() + self.weights_dtype = weights_dtype + self.use_quantized_matmul = use_quantized_matmul + self.result_dtype = result_dtype + self.result_shape = result_shape + self.register_buffer("scale", scale) + + def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: + return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]) + + def forward(self, weight, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument + return decompress_symmetric(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul) + + +class PackedINTAsymmetricWeightsDecompressor(torch.nn.Module): + def __init__( + self, + scale: torch.Tensor, + zero_point: torch.Tensor, + compressed_weight_shape: torch.Size, + result_dtype: torch.dtype, + result_shape: torch.Size, + weights_dtype: str, + **kwargs, # pylint: disable=unused-argument + ): + super().__init__() + self.weights_dtype = weights_dtype + self.use_quantized_matmul = False + self.compressed_weight_shape = compressed_weight_shape + self.result_dtype = result_dtype + self.result_shape = result_shape + self.register_buffer("scale", scale) + self.register_buffer("zero_point", zero_point) + + def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: + return packed_int_function_dict[self.weights_dtype]["pack"](weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])) + + def forward(self, weight, **kwargs): # pylint: disable=unused-argument + return decompress_packed_int_asymmetric(weight, self.scale, self.zero_point, self.compressed_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype) + + +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, # pylint: disable=unused-argument + ): + super().__init__() + self.weights_dtype = weights_dtype + self.use_quantized_matmul = use_quantized_matmul + self.compressed_weight_shape = compressed_weight_shape + self.result_dtype = result_dtype + self.result_shape = result_shape + self.register_buffer("scale", scale) + + def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: + return pack_int_symetric(weight, self.weights_dtype) + + def forward(self, weight, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument + return decompress_packed_int_symmetric(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, + "int6": PackedINTSymmetricWeightsDecompressor, + "int5": PackedINTSymmetricWeightsDecompressor, + "int4": PackedINTSymmetricWeightsDecompressor, + "int3": PackedINTSymmetricWeightsDecompressor, + "int2": PackedINTSymmetricWeightsDecompressor, + "uint8": AsymmetricWeightsDecompressor, + "uint6": PackedINTAsymmetricWeightsDecompressor, + "uint5": PackedINTAsymmetricWeightsDecompressor, + "uint4": PackedINTAsymmetricWeightsDecompressor, + "uint3": PackedINTAsymmetricWeightsDecompressor, + "uint2": PackedINTAsymmetricWeightsDecompressor, + "uint1": AsymmetricWeightsDecompressor, + "bool": AsymmetricWeightsDecompressor, + "float8_e4m3fn": SymmetricWeightsDecompressor, + "float8_e4m3fnuz": SymmetricWeightsDecompressor, + "float8_e5m2": SymmetricWeightsDecompressor, + "float8_e5m2fnuz": SymmetricWeightsDecompressor, +} + + +if shared.opts.sdnq_decompress_compile: + try: + torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) + decompress_asymmetric = torch.compile(decompress_asymmetric, fullgraph=True) + decompress_symmetric = torch.compile(decompress_symmetric, fullgraph=True) + decompress_packed_int_asymmetric = torch.compile(decompress_packed_int_asymmetric, fullgraph=True) + decompress_packed_int_symmetric = torch.compile(decompress_packed_int_symmetric, fullgraph=True) + except Exception as e: + shared.log.warning(f"Quantization: type=sdnq Decompress using torch.compile is not available: {e}") diff --git a/modules/sdnq/forward.py b/modules/sdnq/forward.py new file mode 100644 index 000000000..e5f345682 --- /dev/null +++ b/modules/sdnq/forward.py @@ -0,0 +1,402 @@ +# pylint: disable=redefined-builtin,no-member,protected-access + +from typing import Callable, List, Tuple, Optional +import torch +from modules import shared + +from .common import conv_types, conv_transpose_types +from .decompressor import decompress_symmetric +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: + return quantized_conv_forward_int8_matmul + else: + if use_tensorwise_fp8_matmul: + return quantized_conv_forward_fp8_matmul_tensorwise + else: + return quantized_conv_forward_fp8_matmul + else: + return quantized_conv_forward + elif layer_class_name in conv_transpose_types: + if layer_class_name.endswith("1d"): + return quantized_conv_transpose_1d_forward + elif layer_class_name.endswith("2d"): + return quantized_conv_transpose_2d_forward + elif layer_class_name.endswith("3d"): + return quantized_conv_transpose_3d_forward + else: + if use_quantized_matmul: + if is_integer: + return quantized_linear_forward_int8_matmul + else: + if use_tensorwise_fp8_matmul: + return quantized_linear_forward_fp8_matmul_tensorwise + else: + return quantized_linear_forward_fp8_matmul + else: + 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.div(input.abs().amax(dim=-1, keepdims=True), 448) + input = torch.div(input, input_scale).clamp_(-448, 448).to(torch.float8_e4m3fn) + input_scale = input_scale.to(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.div(input.abs().amax(dim=-1, keepdims=True), 448) + input = torch.div(input, input_scale).clamp_(-448, 448).to(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.div(input.abs().amax(dim=-1, keepdims=True), 127) + input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(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) + result = decompress_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) + if bias is not None: + result.add_(bias) + return result + + +def int8_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + 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_int_symetric(weight, compressed_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) + result = decompress_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape) + if bias is not None: + result.add_(bias) + return result + + +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 = [] + 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 bias is not None: + result.add_(bias) + + 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 = decompress_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, mm_output_shape) + 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 = decompress_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape) + if bias is not None: + result.add_(bias) + + 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, + compressed_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 compressed_weight_shape is not None: + weight = unpack_int_symetric(weight, compressed_weight_shape, weights_dtype, dtype=torch.int8, transpose=True) + + if groups == 1: + result = decompress_symmetric(torch._int_mm(input, weight), scale, return_dtype, mm_output_shape) + 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 = decompress_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape) + if bias is not None: + result.add_(bias) + + 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_decompressor(self.weight, skip_quantized_matmul=True), self.bias) + return fp8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.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_decompressor(self.weight, skip_quantized_matmul=True), self.bias) + return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_decompressor.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_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), self.sdnq_decompressor.weights_dtype) + + +def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor: + return torch.nn.functional.linear(input, self.sdnq_decompressor(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_decompressor(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_decompressor.scale, + self.sdnq_decompressor.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_decompressor(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_decompressor.scale, + self.sdnq_decompressor.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_decompressor(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_decompressor.scale, + self.sdnq_decompressor.result_shape, + getattr(self.sdnq_decompressor, "compressed_weight_shape", None), + self.sdnq_decompressor.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_decompressor(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_decompressor(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_decompressor(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_decompressor(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) + + +if shared.opts.sdnq_decompress_compile: + try: + torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) + int8_matmul = torch.compile(int8_matmul, fullgraph=True) + fp8_matmul = torch.compile(fp8_matmul, fullgraph=True) + fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True) + conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True) + conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True) + conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True) + except Exception as e: + shared.log.warning(f"Quantization: type=sdnq MatMul using torch.compile is not available: {e}") diff --git a/modules/sdnq/packed_int.py b/modules/sdnq/packed_int.py new file mode 100644 index 000000000..351dc7546 --- /dev/null +++ b/modules/sdnq/packed_int.py @@ -0,0 +1,212 @@ +# pylint: disable=redefined-builtin,no-member,protected-access + +from typing import Optional +import torch + +from .common import dtype_dict + + +def pack_int_symetric(tensor: torch.ByteTensor, weights_dtype: str) -> torch.ByteTensor: + return packed_int_function_dict[weights_dtype]["pack"](tensor.to(dtype=dtype_dict[weights_dtype]["torch_dtype"]).sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"])) + + +def unpack_int_symetric(packed_tensor: torch.CharTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, transpose: Optional[bool] = False) -> torch.ByteTensor: + if dtype is None: + dtype = dtype_dict[weights_dtype]["torch_dtype"] + result = packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"]) + if transpose: + result = result.transpose(0,1) + return result + + +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(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 2), 192)), + torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 3], 4), 192)), + torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 6)), + ), + dim=-1 + ) + return packed_tensor + + +def pack_uint5(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_left_shift(packed_tensor[:, 5], 5)), + torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_left_shift(packed_tensor[:, 6], 5)), + torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 7], 5)), + torch.bitwise_or( + packed_tensor[:, 3], + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 5], 2), 96), + torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128), + ), + ), + torch.bitwise_or( + packed_tensor[:, 4], + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 2), 96), + torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128), + ), + ), + ), + dim=-1 + ) + return packed_tensor + + +def pack_uint4(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, 2) + packed_tensor = torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 4)) + return packed_tensor + + +def pack_uint3(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( + torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 3)), + torch.bitwise_left_shift(packed_tensor[:, 6], 6), + ), + torch.bitwise_or( + torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 3)), + torch.bitwise_left_shift(packed_tensor[:, 7], 6), + ), + torch.bitwise_or( + torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_left_shift(packed_tensor[:, 5], 3)), + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 4), 64), + torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128), + ) + ), + ), + dim=-1 + ) + return packed_tensor + + +def pack_uint2(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.bitwise_or( + torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 2)), + torch.bitwise_or(torch.bitwise_left_shift(packed_tensor[:, 2], 4), torch.bitwise_left_shift(packed_tensor[:, 3], 6)), + ) + 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_uint5(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: + result = torch.stack( + ( + torch.bitwise_and(packed_tensor[:, 0], 31), + torch.bitwise_and(packed_tensor[:, 1], 31), + torch.bitwise_and(packed_tensor[:, 2], 31), + torch.bitwise_and(packed_tensor[:, 3], 31), + torch.bitwise_and(packed_tensor[:, 4], 31), + torch.bitwise_or( + torch.bitwise_right_shift(packed_tensor[:, 0], 5), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 2), 24), + ), + torch.bitwise_or( + torch.bitwise_right_shift(packed_tensor[:, 1], 5), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 2), 24), + ), + torch.bitwise_or( + torch.bitwise_right_shift(packed_tensor[:, 2], 5), + torch.bitwise_or( + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 3), 16), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 4), 8), + ), + ), + ), + dim=-1 + ).reshape(shape) + 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 + + +def unpack_uint3(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: + result = torch.stack( + ( + torch.bitwise_and(packed_tensor[:, 0], 7), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 3), 7), + torch.bitwise_and(packed_tensor[:, 1], 7), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 3), 7), + torch.bitwise_and(packed_tensor[:, 2], 7), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 7), + torch.bitwise_or( + torch.bitwise_right_shift(packed_tensor[:, 0], 6), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 4), 4), + ), + torch.bitwise_or( + torch.bitwise_right_shift(packed_tensor[:, 1], 6), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 5), 4), + ), + ), + dim=-1 + ).reshape(shape) + return result + + +def unpack_uint2(packed_tensor: torch.Tensor, shape: torch.Size) -> torch.Tensor: + result = torch.stack( + ( + torch.bitwise_and(packed_tensor, 3), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 2), 3), + torch.bitwise_and(torch.bitwise_right_shift(packed_tensor, 4), 3), + torch.bitwise_right_shift(packed_tensor, 6), + ), + dim=-1 + ).reshape(shape) + return result + + +packed_int_function_dict = { + "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}, + "uint6": {"pack": pack_uint6, "unpack": unpack_uint6}, + "uint5": {"pack": pack_uint5, "unpack": unpack_uint5}, + "uint4": {"pack": pack_uint4, "unpack": unpack_uint4}, + "uint3": {"pack": pack_uint3, "unpack": unpack_uint3}, + "uint2": {"pack": pack_uint2, "unpack": unpack_uint2}, +}