Pre-load support for NNCF

modules/model_quant.py

@@ -104,11 +104,35 @@ def create_quanto_config(kwargs = None, allow_quanto: bool = True, module: str =
     return kwargs
 
 
+def create_nncf_config(kwargs = None, allow_nncf: bool = True, module: str = 'Model'):
+    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
+            diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["nncf"] = NNCFConfig
+            transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["nncf"] = NNCFConfig
+
+            nncf_config = NNCFConfig(weights_dtype=shared.opts.nncf_compress_weights_mode.lower())
+            log.debug(f'Quantization: module="{module}" type=nncf dtype={shared.opts.nncf_compress_weights_mode}')
+            if kwargs is None:
+                return nncf_config
+            else:
+                kwargs['quantization_config'] = nncf_config
+                return kwargs
+    return kwargs
+
+
 def check_quant(module: str = ''):
     from modules import shared
-    if 'Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization:
+    if 'Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization or 'Model' in shared.opts.nncf_compress_weights:
         return True
-    if module in shared.opts.bnb_quantization or module in shared.opts.torchao_quantization or module in shared.opts.quanto_quantization:
+    if module in shared.opts.bnb_quantization or module in shared.opts.torchao_quantization or module in shared.opts.quanto_quantization or module in shared.opts.nncf_compress_weights:
         return True
     return False
 
@@ -142,6 +166,11 @@ def create_config(kwargs = None, allow: bool = True, module: str = 'Model'):
         if debug:
             log.trace(f'Quantization: type=quanto config={kwargs.get("quantization_config", None)}')
         return kwargs
+    kwargs = create_nncf_config(kwargs, allow_nncf=allow, module=module)
+    if kwargs is not None and 'quantization_config' in kwargs:
+        if debug:
+            log.trace(f'Quantization: type=nncf config={kwargs.get("quantization_config", None)}')
+        return kwargs
     return kwargs
 
 
@@ -299,16 +328,25 @@ def nncf_send_to_device(model, device):
         nncf_send_to_device(child, device)
 
 
-def nncf_compress_model(model, op=None, sd_model=None):
+def nncf_compress_model(model, op=None, sd_model=None, send_to_device=True, do_gc=True):
     from modules import devices, shared
     global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
     nncf = load_nncf('Quantize model: type=NNCF')
     model.eval()
+    if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
+        import torch
+        from modules.model_quant_nncf import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
+        for i in range(len(model.encoder.block)):
+            model.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense(
+                model.encoder.block[i].layer[1].DenseReluDense,
+                dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
+            )
     backup_embeddings = None
     if hasattr(model, "get_input_embeddings"):
         backup_embeddings = copy.deepcopy(model.get_input_embeddings())
     model = nncf.compress_weights(model)
-    nncf_send_to_device(model, devices.device)
+    if send_to_device:
+        nncf_send_to_device(model, devices.device)
     if hasattr(model, "set_input_embeddings") and backup_embeddings is not None:
         model.set_input_embeddings(backup_embeddings)
     if op is not None and shared.opts.nncf_quantize_shuffle_weights:
@@ -318,7 +356,8 @@ def nncf_compress_model(model, op=None, sd_model=None):
                 getattr(getattr(sd_model, last_model_names[0]), last_model_names[1]).to(quant_last_model_device)
             else:
                 getattr(sd_model, quant_last_model_name).to(quant_last_model_device)
-            devices.torch_gc(force=True)
+            if do_gc:
+                devices.torch_gc(force=True)
         if shared.cmd_opts.medvram or shared.cmd_opts.lowvram or shared.opts.diffusers_offload_mode != "none":
             quant_last_model_name = op
             quant_last_model_device = model.device
@@ -326,7 +365,8 @@ def nncf_compress_model(model, op=None, sd_model=None):
             quant_last_model_name = None
             quant_last_model_device = None
         model.to(devices.device)
-    devices.torch_gc(force=True)
+    if do_gc:
+        devices.torch_gc(force=True)
     return model
 
 
@@ -512,7 +552,7 @@ def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, al
 
 def do_post_load_quant(sd_model):
     from modules import shared
-    if shared.opts.nncf_compress_weights and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
+    if shared.opts.nncf_compress_weights and shared.opts.nncf_compress_mode == 'post' and not (shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"):
         sd_model = nncf_compress_weights(sd_model)
     if shared.opts.optimum_quanto_weights:
         sd_model = optimum_quanto_weights(sd_model)

modules/model_quant_nncf.py

@@ -0,0 +1,219 @@
+from typing import Any, Dict, List, 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 accelerate import init_empty_weights
+
+from modules import devices
+
+
+class NNCF_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
+
+
+class QuantizationMethod(str, Enum):
+    NNCF = "nncf"
+
+
+class ConvertToModule(torch.nn.Module):
+    def __init__(self, module):
+        super().__init__()
+        self.module = module
+
+    def forward(self, *args, **kwargs):
+        return self.module(*args, **kwargs)
+
+
+class NNCFQuantizer(DiffusersQuantizer):
+    r"""
+    Diffusers Quantizer for NNCF
+    """
+
+    requires_parameters_quantization = True
+    use_keep_in_fp32_modules = True
+    requires_calibration = False
+    required_packages = ["nncf"]
+
+    def __init__(self, quantization_config, **kwargs):
+        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,
+    ):
+        module, tensor_name = get_module_from_name(model, param_name)
+        return module.__class__.__name__.startswith("NNCF") and param_name.endswith(".weight")
+
+    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(
+        self,
+        model,
+        param_value: "torch.Tensor",
+        param_name: str,
+        target_device: "torch.device",
+        state_dict: Dict[str, Any],
+        unexpected_keys: List[str],
+        **kwargs,
+    ):
+        # load the model params to target_device first
+        layer, tensor_name = get_module_from_name(model, param_name)
+        layer._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device)
+
+        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):
+            from nncf.torch.quantization.quantize_functions import get_scale_zp_from_input_low_input_high
+            from nncf.torch.quantization.weights_compression import WeightsDecompressor
+            from nncf.torch.layers import NNCFEmbedding
+
+            if not isinstance(layer, torch.nn.Embedding) and not isinstance(layer, NNCFEmbedding):
+                target_dim = layer.target_weight_dim_for_compression
+                stat_dim = (target_dim + 1) % 2
+                input_low = torch.min(layer.weight, dim=stat_dim).values.detach()
+                input_high = torch.max(layer.weight, dim=stat_dim).values.detach()
+                scale, zero_point = get_scale_zp_from_input_low_input_high(0, 255, input_low, input_high)
+
+                scale = scale.unsqueeze(stat_dim)
+                zero_point = zero_point.unsqueeze(stat_dim)
+                layer.register_pre_forward_operation(WeightsDecompressor(zero_point, scale))
+
+                compressed_weight = layer.weight.data / scale + zero_point
+                compressed_weight = torch.clamp(torch.round(compressed_weight), 0, 255)
+
+                layer.weight.requires_grad = False
+                layer.weight.data = compressed_weight.type(dtype=torch.uint8)
+
+    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":
+        return torch.uint8
+
+    def _process_model_before_weight_loading(
+        self,
+        model,
+        device_map,
+        keep_in_fp32_modules: List[str] = [],
+        **kwargs,
+    ):
+        from nncf.torch.nncf_module_replacement import replace_modules_by_nncf_modules
+
+        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]
+
+        self.modules_to_not_convert.extend(keep_in_fp32_modules)
+        model.config.quantization_config = self.quantization_config
+
+        if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
+            for i in range(len(model.encoder.block)):
+                model.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense(
+                    model.encoder.block[i].layer[1].DenseReluDense,
+                    dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
+                )
+
+        with init_empty_weights():
+            model, _ = replace_modules_by_nncf_modules(model)
+
+    def _process_model_after_weight_loading(self, model, **kwargs):
+        from modules.model_quant import nncf_send_to_device
+        nncf_send_to_device(model, devices.device)
+        return model
+
+    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]:
+        """
+        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]:
+        """
+        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]:
+        """
+        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 NNCFConfig(QuantizationConfigMixin):
+    """
+    This is a wrapper class about all possible attributes and features that you can play with a model that has been
+    loaded using `quanto`.
+
+    Args:
+        weights_dtype (`str`, *optional*, defaults to `"int8"`):
+            The target dtype for the weights after quantization. Supported values are ("int8")
+       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__(
+        self,
+        weights_dtype: str = "int8",
+        modules_to_not_convert: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        self.quant_method = QuantizationMethod.NNCF
+        self.weights_dtype = weights_dtype
+        self.modules_to_not_convert = modules_to_not_convert
+
+        self.post_init()
+
+    def post_init(self):
+        r"""
+        Safety checker that arguments are correct
+        """
+        accepted_weights = ["int8", "uint8"]
+        if self.weights_dtype not in accepted_weights:
+            raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}")
+

modules/model_te.py

@@ -69,23 +69,20 @@ def load_t5(name=None, cache_dir=None):
         quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True)
         t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
+    elif 'int8' in name.lower():
+        from modules.model_quant import create_nncf_config
+        quantization_config = create_nncf_config(kwargs=None, allow_nncf=True, module="any")
+        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+
+    elif 'qint4' in name.lower():
+        model_quant.load_quanto('Load model: type=T5')
+        quantization_config = transformers.QuantoConfig(weights='int4')
+        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+
     elif 'qint8' in name.lower():
         model_quant.load_quanto('Load model: type=T5')
-        from modules.model_quant import optimum_quanto_model
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype)
-        t5 = optimum_quanto_model(t5, weights="qint8", activations="none")
-
-    elif 'int8' in name.lower():
-        install('nncf==2.7.0', quiet=True)
-        from modules.model_quant import nncf_compress_model
-        from modules.sd_hijack import NNCF_T5DenseGatedActDense
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype)
-        for i in range(len(t5.encoder.block)):
-            t5.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense(
-                t5.encoder.block[i].layer[1].DenseReluDense,
-                dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
-            )
-        t5 = nncf_compress_model(t5)
+        quantization_config = transformers.QuantoConfig(weights='int8')
+        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
     elif '/' in name:
         shared.log.debug(f'Load model: type=T5 repo={name}')

modules/sd_hijack.py

@@ -283,26 +283,6 @@ class EmbeddingsWithFixes(torch.nn.Module):
         return torch.stack(vecs)
 
 
-class NNCF_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
-
-
 def add_circular_option_to_conv_2d():
     conv2d_constructor = torch.nn.Conv2d.__init__
 

modules/sd_models_utils.py

@@ -171,40 +171,12 @@ def apply_function_to_model(sd_model, function, options, op=None):
             if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model.decoder_pipe, 'text_encoder') and hasattr(sd_model.decoder_pipe.text_encoder, 'config'):
                 sd_model.decoder_pipe.text_encoder = function(sd_model.decoder_pipe.text_encoder, op="decoder_pipe.text_encoder", sd_model=sd_model)
             else:
-                if op == "nncf" and sd_model.text_encoder.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
-                    from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
-                    for i in range(len(sd_model.text_encoder.encoder.block)):
-                        sd_model.text_encoder.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense(
-                            sd_model.text_encoder.encoder.block[i].layer[1].DenseReluDense,
-                            dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
-                        )
                 sd_model.text_encoder = function(sd_model.text_encoder, op="text_encoder", sd_model=sd_model)
         if hasattr(sd_model, 'text_encoder_2') and hasattr(sd_model.text_encoder_2, 'config'):
-            if op == "nncf" and sd_model.text_encoder_2.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
-                from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
-                for i in range(len(sd_model.text_encoder_2.encoder.block)):
-                    sd_model.text_encoder_2.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense(
-                        sd_model.text_encoder_2.encoder.block[i].layer[1].DenseReluDense,
-                        dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
-                    )
             sd_model.text_encoder_2 = function(sd_model.text_encoder_2, op="text_encoder_2", sd_model=sd_model)
         if hasattr(sd_model, 'text_encoder_3') and hasattr(sd_model.text_encoder_3, 'config'):
-            if op == "nncf" and sd_model.text_encoder_3.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
-                from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
-                for i in range(len(sd_model.text_encoder_3.encoder.block)):
-                    sd_model.text_encoder_3.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense(
-                        sd_model.text_encoder_3.encoder.block[i].layer[1].DenseReluDense,
-                        dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
-                    )
             sd_model.text_encoder_3 = function(sd_model.text_encoder_3, op="text_encoder_3", sd_model=sd_model)
         if hasattr(sd_model, 'text_encoder_4') and hasattr(sd_model.text_encoder_4, 'config'):
-            if op == "nncf" and sd_model.text_encoder_4.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
-                from modules.sd_hijack import NNCF_T5DenseGatedActDense # T5DenseGatedActDense uses fp32
-                for i in range(len(sd_model.text_encoder_4.encoder.block)):
-                    sd_model.text_encoder_4.encoder.block[i].layer[1].DenseReluDense = NNCF_T5DenseGatedActDense(
-                        sd_model.text_encoder_4.encoder.block[i].layer[1].DenseReluDense,
-                        dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
-                    )
             sd_model.text_encoder_4 = function(sd_model.text_encoder_4, op="text_encoder_4", sd_model=sd_model)
         if hasattr(sd_model, 'prior_pipe') and hasattr(sd_model.prior_pipe, 'text_encoder') and hasattr(sd_model.prior_pipe.text_encoder, 'config'):
             sd_model.prior_pipe.text_encoder = function(sd_model.prior_pipe.text_encoder, op="prior_pipe.text_encoder", sd_model=sd_model)

modules/shared.py

@@ -537,6 +537,7 @@ options_templates.update(options_section(('quantization', "Quantization Settings
 
     "nncf_compress_sep": OptionInfo("<h2>NNCF: Neural Network Compression Framework</h2>", "", gr.HTML),
     "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
+    "nncf_compress_mode": OptionInfo("post", "Quantization mode", gr.Dropdown, {"choices": ['pre', 'post'], "visible": native}),
     "nncf_compress_weights_mode": OptionInfo("INT8", "Quantization type", gr.Dropdown, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'] if cmd_opts.use_openvino else ['INT8']}),
     "nncf_compress_weights_raito": OptionInfo(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}),
     "nncf_compress_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 512, "step": 1, "visible": cmd_opts.use_openvino}),