From 2d79380bd71e376d9b1460248c5583ea504b5efc Mon Sep 17 00:00:00 2001 From: Disty0 Date: Mon, 26 May 2025 01:12:28 +0300 Subject: [PATCH] NNCF implement better layer hijacks and remove all NNCF imports --- installer.py | 3 - modules/lora/lora_apply.py | 21 +++--- modules/model_quant.py | 40 ---------- modules/model_quant_nncf.py | 147 +++++++++++++++++++----------------- 4 files changed, 87 insertions(+), 124 deletions(-) diff --git a/installer.py b/installer.py index 714db5ff2..769256404 100644 --- a/installer.py +++ b/installer.py @@ -787,8 +787,6 @@ def install_torch_addons(): install('DeepCache') if opts.get('cuda_compile_backend', '') == 'olive-ai': install('olive-ai') - if opts.get('nncf_compress_weights', False) and not args.use_openvino: - install('nncf==2.16.0', 'nncf') if opts.get('optimum_quanto_weights', False): install('optimum-quanto==0.2.7', 'optimum-quanto') if opts.get('torchao_quantization', False): @@ -1175,7 +1173,6 @@ def install_optional(): install('albumentations==1.4.3', ignore=True) install('pydantic==1.10.21', ignore=True) reload('pydantic', '1.10.21') - install('nncf==2.16.0', ignore=True) install('gguf', ignore=True) install('av', ignore=True) try: diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index 68f723188..7b00ee389 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -45,8 +45,8 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n self.network_weights_backup = True else: self.network_weights_backup = weight.clone().to(devices.cpu) - if self.__class__.__name__.startswith('NNCF') and hasattr(self, "pre_ops") and len(self.pre_ops) == 1: - self.nncf_decompressor_backup = self.pre_ops["0"].to(devices.cpu) + if hasattr(self, "nncf_decompressor"): + self.nncf_decompressor_backup = self.nncf_decompressor.to(devices.cpu) if bias_backup is None: if getattr(self, 'bias', None) is not None: @@ -79,10 +79,10 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn. continue try: t0 = time.time() - if self.__class__.__name__.startswith('NNCF') and hasattr(self, "pre_ops") and len(self.pre_ops) == 1: + if hasattr(self, "nncf_decompressor"): return_device = self.weight.data.device self.weight.data = self.weight.data.to(devices.device) - weight = self.pre_ops["0"].to(devices.device)(self, return_decompressed_only=True) + weight = self.nncf_decompressor.to(devices.device)(self, return_decompressed_only=True) self.weight.data = self.weight.data.to(return_device) else: weight = self.weight.to(devices.device) # must perform calc on gpu due to performance @@ -139,19 +139,18 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G # weight._quantize(devices.device) / weight.to(device=device) except Exception as e: 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 self.__class__.__name__.startswith('NNCF') and hasattr(self, "pre_ops") and len(self.pre_ops) == 1: + elif not bias and hasattr(self, "nncf_decompressor"): num_bits = None is_asym_mode = None try: from modules.model_quant_nncf import nncf_compress_layer - num_bits = self.pre_ops["0"].num_bits - is_asym_mode = self.pre_ops["0"].quantization_mode == "asymmetric" + num_bits = self.nncf_decompressor.num_bits + is_asym_mode = self.nncf_decompressor.quantization_mode == "asymmetric" self.weight = torch.nn.Parameter(model_weights.to(devices.device), requires_grad=False) - dequant_weight = self.pre_ops["0"](self, return_decompressed_only=True) + dequant_weight = self.nncf_decompressor(self, return_decompressed_only=True) new_weight = dequant_weight.to(devices.device, dtype=torch.float32) + lora_weights.to(devices.device, dtype=torch.float32) self.weight = torch.nn.Parameter(new_weight, requires_grad=False) - self.pre_ops.pop("0") - self._custom_forward_fn = None # pylint: disable=protected-access + self.nncf_decompressor = None self = nncf_compress_layer(self, num_bits, is_asym_mode, torch_dtype=devices.dtype, quant_conv=shared.opts.nncf_quantize_conv_layers, group_size=shared.opts.nncf_compress_weights_group_size, use_int8_matmul=shared.opts.nncf_decompress_int8_matmul) self = self.to(device) del dequant_weight @@ -219,7 +218,7 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn else: self.weight = torch.nn.Parameter(weights_backup.to(device), requires_grad=False) if hasattr(self, "nncf_decompressor_backup"): - self.pre_ops["0"] = self.nncf_decompressor_backup.to(device) + self.nncf_decompressor = self.nncf_decompressor_backup.to(device) if bias_backup is not None: self.bias = None diff --git a/modules/model_quant.py b/modules/model_quant.py index a63123b34..3b453691d 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -9,7 +9,6 @@ from installer import installed, install, log, setup_logging ao = None bnb = None -intel_nncf = None optimum_quanto = None quant_last_model_name = None quant_last_model_device = None @@ -108,9 +107,6 @@ def create_nncf_config(kwargs = None, allow_nncf: bool = True, module: str = 'Mo from modules import shared if len(shared.opts.nncf_compress_weights) > 0 and (shared.opts.nncf_compress_mode == 'pre') and allow_nncf: if 'Model' in shared.opts.nncf_compress_weights or (module is not None and module in shared.opts.nncf_compress_weights) or module == 'any': - load_nncf(silent=True) - if intel_nncf is None: - return kwargs from modules.model_quant_nncf import NNCFQuantizer, NNCFConfig diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["nncf"] = NNCFQuantizer transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["nncf"] = NNCFQuantizer @@ -259,35 +255,6 @@ def load_quanto(msg='', silent=False): return None -def load_nncf(msg='', silent=False): - global intel_nncf # pylint: disable=global-statement - if intel_nncf is not None: - return intel_nncf - if not installed('nncf'): - install('nncf==2.16.0', quiet=True) - log.warning('Quantization: nncf installed please restart') - install('jstyleson', quiet=True) - install('texttable', quiet=True) - install('tabulate', quiet=True) - try: - import nncf - intel_nncf = nncf - try: - nncf.common.logging.logger.warn_bkc_version_mismatch = lambda *args, **kwargs: None # silence the pytorch version warning - except Exception: - pass - fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access - log.debug(f'Quantization: type=nncf version={nncf.__version__} fn={fn}') # pylint: disable=protected-access - return intel_nncf - except Exception as e: - if len(msg) > 0: - log.error(f"{msg} failed to import nncf: {e}") - intel_nncf = None - if not silent: - raise - return None - - def apply_layerwise(sd_model, quiet:bool=False): import torch from diffusers.quantizers import quantization_config @@ -334,11 +301,7 @@ def apply_layerwise(sd_model, quiet:bool=False): def nncf_compress_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 accelerate import init_empty_weights - - load_nncf('Quantize model: type=NNCF') from modules.model_quant_nncf import apply_nncf_to_module - from nncf.torch.nncf_module_replacement import replace_modules_by_nncf_modules # get around lazy import model.eval() @@ -355,9 +318,6 @@ def nncf_compress_model(model, op=None, sd_model=None, do_gc=True): if hasattr(model, "get_input_embeddings"): backup_embeddings = copy.deepcopy(model.get_input_embeddings()) - with init_empty_weights(): - model, _ = replace_modules_by_nncf_modules(model) - num_bits = 8 if shared.opts.nncf_compress_weights_mode in {"INT8", "INT8_SYM", "INT8_ASYM"} else 4 is_asym_mode = shared.opts.nncf_compress_weights_mode in {"INT8", "INT4", "INT8_ASYM", "INT4_ASYM"} model = apply_nncf_to_module(model, num_bits, is_asym_mode, quant_conv=shared.opts.nncf_quantize_conv_layers) diff --git a/modules/model_quant_nncf.py b/modules/model_quant_nncf.py index 3d8fb4624..2e07eced8 100644 --- a/modules/model_quant_nncf.py +++ b/modules/model_quant_nncf.py @@ -8,7 +8,6 @@ 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 import init_empty_weights from accelerate.utils import CustomDtype from modules import devices, shared @@ -42,22 +41,29 @@ class QuantizationMethod(str, Enum): def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_conv=False, group_size=0, use_int8_matmul=False, param_name=None): # pylint: disable=unused-argument - if layer.__class__.__name__ in allowed_types: + 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 - result_shape = None - if layer.__class__.__name__ in conv_types: - if is_asym_mode or not quant_conv: # don't quant convs with asym mode + if layer_class_name in conv_types: + if not quant_conv: return layer reduction_axes = [i for i in range(layer.weight.ndim) if i != 0] use_int8_matmul = False - if layer.__class__.__name__ in conv_transpose_types: - if is_asym_mode or not quant_conv: # don't quant convs with asym mode + is_conv_type = True + elif layer_class_name in conv_transpose_types: + if not quant_conv: return layer reduction_axes = [i for i in range(layer.weight.ndim) if i != 1] use_int8_matmul = False + is_conv_transpose_type = True else: + is_linear_type = True reduction_axes = -1 channel_size = layer.weight.shape[-1] use_int8_matmul = use_int8_matmul and not is_asym_mode and channel_size >= 32 and layer.weight.shape[0] >= 32 @@ -106,12 +112,9 @@ def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_c zero_point = zero_point.to(torch_dtype) if use_int8_matmul: - layer._custom_forward_fn = linear_forward_int8_matmul # pylint: disable=protected-access scale = scale.squeeze(-1) if num_bits == 8: compressed_weight = compressed_weight.transpose(0,1) - else: - layer._custom_forward_fn = None # pylint: disable=protected-access if num_bits == 4: if is_asym_mode: @@ -128,7 +131,6 @@ def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_c compressed_weight_shape=compressed_weight.shape, result_dtype=torch_dtype, result_shape=result_shape, - use_int8_matmul=use_int8_matmul, ) else: if is_asym_mode: @@ -143,14 +145,30 @@ def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_c scale=scale.data, result_dtype=torch_dtype, result_shape=result_shape, - use_int8_matmul=use_int8_matmul, ) compressed_weight = decompressor.pack_weight(compressed_weight).to(return_device) decompressor = decompressor.to(return_device) - layer.register_pre_forward_operation(decompressor) + layer.weight.requires_grad = False layer.weight.data = compressed_weight + layer.nncf_decompressor = decompressor + + if is_linear_type: + if use_int8_matmul: + layer.forward = quantized_linear_forward_int8_matmul + else: + layer.forward = quantized_linear_forward + elif is_conv_type: + 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__) return layer @@ -159,7 +177,7 @@ def apply_nncf_to_module(model, num_bits, is_asym_mode, quant_conv=False): if not has_children: return model for param_name, module in model.named_children(): - if module.__class__.__name__.startswith("NNCF") and hasattr(module, "weight") and module.weight is not None: + if hasattr(module, "weight") and module.weight is not None: module = nncf_compress_layer( module, num_bits, @@ -205,8 +223,7 @@ class NNCFQuantizer(DiffusersQuantizer): state_dict: Dict[str, Any], **kwargs, ): - module, _ = get_module_from_name(model, param_name) - return module.__class__.__name__.startswith("NNCF") and param_name.endswith(".weight") + return param_name.endswith(".weight") def check_quantized_param(self, *args, **kwargs) -> bool: """ @@ -228,12 +245,6 @@ class NNCFQuantizer(DiffusersQuantizer): layer, tensor_name = get_module_from_name(model, param_name) layer._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device) # pylint: disable=protected-access - # nncf_padding_value somehow ends up in the meta device with cogvideo even if we don't use init_empty_weights - # set it to the default value if it is in the meta device: - if layer.__class__.__name__ == "NNCFConv2d" and hasattr(layer, "get_padding_value_ref") and hasattr(layer, "_set_padding_value"): - if layer.get_padding_value_ref().device == torch.device("meta"): - layer._set_padding_value(torch.zeros([1])) - 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 = nncf_compress_layer( @@ -266,18 +277,13 @@ class NNCFQuantizer(DiffusersQuantizer): keep_in_fp32_modules: List[str] = [], **kwargs, ): - from nncf.torch.nncf_module_replacement import replace_modules_by_nncf_modules - + 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) - model.config.quantization_config = self.quantization_config - with init_empty_weights(): - model, _ = replace_modules_by_nncf_modules(model) - def _process_model_after_weight_loading(self, model, **kwargs): return model @@ -470,28 +476,51 @@ def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTenso def int8_matmul( - input: torch.Tensor, + input: torch.FloatTensor, weight: torch.Tensor, - scale: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, compressed_weight_shape: torch.Size, -): +) -> torch.FloatTensor: if compressed_weight_shape is not None: weight = unpack_int4_compiled(weight, compressed_weight_shape, transpose=True) return_dtype = input.dtype output_shape = list(input.shape) output_shape[-1] = weight.shape[-1] input, scale = quantize_int8_matmul_input_compiled(input, scale) - return decompress_symmetric_compiled(torch._int_mm(input, weight), scale, return_dtype, output_shape) # pylint: disable=protected-access + result = decompress_symmetric_compiled(torch._int_mm(input, weight), scale, return_dtype, output_shape) # pylint: disable=protected-access + if bias is not None: + result.add_(bias) + return result -class linear_forward_int8_matmul(): - def __func__(self, input) -> torch.FloatTensor: - if self.pre_ops["0"].skip_int8_matmul: - return torch.nn.functional.linear(input, self.weight, self.bias) - result = int8_matmul(input, self.weight, self.pre_ops["0"].scale, getattr(self.pre_ops["0"], "compressed_weight_shape", None)) - if self.bias is not None: - result.add_(self.bias) - return result +def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: + if torch.numel(input[0]) / input[0].shape[-1] < 32: + return torch.nn.functional.linear(input, self.nncf_decompressor(self, return_decompressed_only=True, skip_int8_matmul=True), self.bias) + return int8_matmul(input, self.weight, self.bias, self.nncf_decompressor.scale, getattr(self.nncf_decompressor, "compressed_weight_shape", None)) + + +def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor: + return torch.nn.functional.linear(input, self.nncf_decompressor(self, return_decompressed_only=True), self.bias) + + +def quantized_conv_forward(self, input) -> torch.FloatTensor: + return self._conv_forward(input, self.nncf_decompressor(self, return_decompressed_only=True), 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.nncf_decompressor(self, return_decompressed_only=True), 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.nncf_decompressor(self, return_decompressed_only=True), 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.nncf_decompressor(self, return_decompressed_only=True), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) class INT8AsymmetricWeightsDecompressor(torch.nn.Module): @@ -530,7 +559,6 @@ class INT8SymmetricWeightsDecompressor(torch.nn.Module): scale: torch.Tensor, result_dtype: torch.dtype, result_shape: torch.Size, - use_int8_matmul: bool, ): super().__init__() self.num_bits = 8 @@ -538,9 +566,6 @@ class INT8SymmetricWeightsDecompressor(torch.nn.Module): self.scale = scale self.result_dtype = result_dtype self.result_shape = result_shape - self.use_int8_matmul = use_int8_matmul - self.skip_int8_matmul = False - self.input_scale = None def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: if debug: @@ -548,17 +573,10 @@ class INT8SymmetricWeightsDecompressor(torch.nn.Module): raise ValueError("Weight values are not in [-128, 127].") return weight.to(dtype=torch.int8) - def forward(self, x, input=None, *args, return_decompressed_only=False): # pylint: disable=unused-argument,keyword-arg-before-vararg - if self.use_int8_matmul: - if input is not None: - if torch.numel(input[0]) / input[0].shape[-1] < 32: - self.skip_int8_matmul = True - else: - self.skip_int8_matmul = False - return - result = decompress_symmetric_compiled(x.weight.transpose(0,1), self.scale.unsqueeze(-1), self.result_dtype, self.result_shape) - else: - result = decompress_symmetric_compiled(x.weight, self.scale, self.result_dtype, self.result_shape) + def forward(self, x, input=None, *args, return_decompressed_only=False, skip_int8_matmul=False): # pylint: disable=unused-argument,keyword-arg-before-vararg + if skip_int8_matmul: + return decompress_int4_symmetric_compiled(x.weight, self.scale.unsqueeze(-1), self.compressed_weight_shape, self.result_dtype, self.result_shape) + result = decompress_symmetric_compiled(x.weight, self.scale, self.result_dtype, self.result_shape) if return_decompressed_only: return result else: @@ -604,7 +622,6 @@ class INT4SymmetricWeightsDecompressor(torch.nn.Module): compressed_weight_shape: torch.Size, result_dtype: torch.dtype, result_shape: torch.Size, - use_int8_matmul: bool, ): super().__init__() self.num_bits = 4 @@ -613,9 +630,6 @@ class INT4SymmetricWeightsDecompressor(torch.nn.Module): self.compressed_weight_shape = compressed_weight_shape self.result_dtype = result_dtype self.result_shape = result_shape - self.use_int8_matmul = use_int8_matmul - self.skip_int8_matmul = False - self.input_scale = None def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: if debug: @@ -623,17 +637,10 @@ class INT4SymmetricWeightsDecompressor(torch.nn.Module): raise ValueError("Tensor values are not in [-8, 7].") return pack_int4(weight.to(dtype=torch.int8)) - def forward(self, x, input=None, *arg, return_decompressed_only=False): # pylint: disable=keyword-arg-before-vararg,unused-argument - if self.use_int8_matmul: - if input is not None: - if torch.numel(input[0]) / input[0].shape[-1] < 32: - self.skip_int8_matmul = True - else: - self.skip_int8_matmul = False - return - result = decompress_int4_symmetric_compiled(x.weight, self.scale.unsqueeze(-1), self.compressed_weight_shape, self.result_dtype, self.result_shape) - else: - result = decompress_int4_symmetric_compiled(x.weight, self.scale, self.compressed_weight_shape, self.result_dtype, self.result_shape) + def forward(self, x, input=None, *arg, return_decompressed_only=False, skip_int8_matmul=False): # pylint: disable=keyword-arg-before-vararg,unused-argument + if skip_int8_matmul: + return decompress_int4_symmetric_compiled(x.weight, self.scale.unsqueeze(-1), self.compressed_weight_shape, self.result_dtype, self.result_shape) + result = decompress_int4_symmetric_compiled(x.weight, self.scale, self.compressed_weight_shape, self.result_dtype, self.result_shape) if return_decompressed_only: return result else: