NNCF implement better layer hijacks and remove all NNCF imports

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:

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

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)

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: