diff --git a/modules/model_quant.py b/modules/model_quant.py index cc8efcc16..7561da7a6 100644 --- a/modules/model_quant.py +++ b/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) diff --git a/modules/model_quant_nncf.py b/modules/model_quant_nncf.py new file mode 100644 index 000000000..d748c081b --- /dev/null +++ b/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}") + diff --git a/modules/model_te.py b/modules/model_te.py index 698d4506e..d853eaf99 100644 --- a/modules/model_te.py +++ b/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}') diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 57573a493..c62f1057f 100644 --- a/modules/sd_hijack.py +++ b/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__ diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py index 273d1b39d..0a9ca7458 100644 --- a/modules/sd_models_utils.py +++ b/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) diff --git a/modules/shared.py b/modules/shared.py index b7bb69b86..5ba516c62 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -537,6 +537,7 @@ options_templates.update(options_section(('quantization', "Quantization Settings "nncf_compress_sep": OptionInfo("

NNCF: Neural Network Compression Framework

", "", 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}),