From f1d8543cae07d579647edba49c7c0bb284d80f07 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 23 Apr 2025 15:44:09 +0300 Subject: [PATCH] NNCF lora support --- CHANGELOG.md | 1 + modules/lora/lora_apply.py | 31 ++++++++++++++++++-- modules/model_quant_nncf.py | 56 +++++++++++++++++++++++++++++-------- 3 files changed, 74 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 73e4a5b2a..6b90e6878 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ major refactoring of NNCF quantization code new quant types: `INT8_SYM` (new default), `INT4` and `INT4_SYM` pre-load quantization support + lora support - **HiDream-I1** optimized offloading and prompt-encode caching it now works in 12GB VRAM / 26GB RAM! - **CogView3** and **CogView4** model loader optimizations diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index 485581a4d..15cf90c5f 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -41,7 +41,12 @@ def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.n else: self.network_weights_backup = weight.clone().to(devices.cpu) if not shared.opts.lora_fuse_diffusers else True else: - self.network_weights_backup = weight.clone().to(devices.cpu) if not shared.opts.lora_fuse_diffusers else True + if shared.opts.lora_fuse_diffusers: + 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 bias_backup is None: if getattr(self, 'bias', None) is not None: @@ -74,7 +79,10 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn. continue try: t0 = time.time() - weight = self.weight.to(devices.device) # must perform calc on gpu due to performance + if self.__class__.__name__.startswith('NNCF') and hasattr(self, "pre_ops") and len(self.pre_ops) == 1: + weight = self.pre_ops["0"](self, return_decompressed_only=True).to(devices.device) + else: + weight = self.weight.to(devices.device) # must perform calc on gpu due to performance updown, ex_bias = module.calc_updown(weight) weight = None del weight @@ -128,6 +136,23 @@ 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: + 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" + self.weight = torch.nn.Parameter(model_weights.to(devices.device), requires_grad=False) + dequant_weight = self.pre_ops["0"](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 = nncf_compress_layer(self, num_bits, is_asym_mode, torch_dtype=devices.dtype, quant_conv=shared.opts.nncf_quantize_conv_layers) + self = self.to(device) + del dequant_weight + except Exception as e: + shared.log.error(f'Network load: type=LoRA quant=nncf cls={self.__class__.__name__} bits={num_bits} is_asym_mode={is_asym_mode} weight={self.weight} lora_weights={lora_weights} {e}') else: try: new_weight = model_weights.to(devices.device) + lora_weights.to(devices.device) @@ -189,6 +214,8 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn network_add_weights(self, model_weights=weights_backup, lora_weights=updown, deactivate=deactivate, device=device, bias=False) 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) if bias_backup is not None: self.bias = None diff --git a/modules/model_quant_nncf.py b/modules/model_quant_nncf.py index 62d732020..56831583d 100644 --- a/modules/model_quant_nncf.py +++ b/modules/model_quant_nncf.py @@ -396,9 +396,13 @@ class INT8AsymmetricWeightsDecompressor(torch.nn.Module): self.zero_point = self.pack_weight(zero_point) self.result_dtype = result_dtype + @property + def num_bits(self): + return 8 + @property def quantization_mode(self): - return"asymmetric" + return "asymmetric" def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: if torch.is_floating_point(weight): @@ -409,9 +413,13 @@ class INT8AsymmetricWeightsDecompressor(torch.nn.Module): raise ValueError(msg) return weight.type(torch.uint8) - def forward(self, x, *args): + def forward(self, x, *args, return_decompressed_only=False): result = decompress_asymmetric(x.weight, self.scale, self.zero_point) - x.weight = result.type(self.result_dtype) + result = result.type(self.result_dtype) + if return_decompressed_only: + return result + else: + x.weight = result class INT8SymmetricWeightsDecompressor(torch.nn.Module): @@ -420,9 +428,13 @@ class INT8SymmetricWeightsDecompressor(torch.nn.Module): self.scale = scale self.result_dtype = result_dtype + @property + def num_bits(self): + return 8 + @property def quantization_mode(self): - return"symmetric" + return "symmetric" def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: if torch.any((weight < -128) | (weight > 127)): @@ -430,9 +442,13 @@ class INT8SymmetricWeightsDecompressor(torch.nn.Module): raise ValueError(msg) return weight.type(torch.int8) - def forward(self, x, *args): + def forward(self, x, *args, return_decompressed_only=False): result = decompress_symmetric(x.weight, self.scale) - x.weight = result.type(self.result_dtype) + result = result.type(self.result_dtype) + if return_decompressed_only: + return result + else: + x.weight = result class INT4AsymmetricWeightsDecompressor(torch.nn.Module): @@ -454,9 +470,13 @@ class INT4AsymmetricWeightsDecompressor(torch.nn.Module): self.result_shape = result_shape self.result_dtype = result_dtype + @property + def num_bits(self): + return 4 + @property def quantization_mode(self): - return"asymmetric" + return "asymmetric" def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: if torch.any((weight < 0) | (weight > 15)): @@ -464,7 +484,7 @@ class INT4AsymmetricWeightsDecompressor(torch.nn.Module): raise ValueError(msg) return pack_uint4(weight.type(torch.uint8)) - def forward(self, x, *args): + def forward(self, x, *args, return_decompressed_only=False): result = unpack_uint4(x.weight) result = result.reshape(self.compressed_weight_shape) @@ -473,7 +493,11 @@ class INT4AsymmetricWeightsDecompressor(torch.nn.Module): result = decompress_asymmetric(result, self.scale, zero_point) result = result.reshape(self.result_shape) if self.result_shape is not None else result - x.weight = result.type(self.result_dtype) + result = result.type(self.result_dtype) + if return_decompressed_only: + return result + else: + x.weight = result class INT4SymmetricWeightsDecompressor(torch.nn.Module): @@ -491,9 +515,13 @@ class INT4SymmetricWeightsDecompressor(torch.nn.Module): self.result_shape = result_shape self.result_dtype = result_dtype + @property + def num_bits(self): + return 4 + @property def quantization_mode(self): - return"symmetric" + return "symmetric" def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: if torch.is_floating_point(weight): @@ -504,10 +532,14 @@ class INT4SymmetricWeightsDecompressor(torch.nn.Module): raise ValueError(msg) return pack_int4(weight.type(torch.int8)) - def forward(self, x, *args): + def forward(self, x, *arg, return_decompressed_only=False): result = unpack_int4(x.weight) result = result.reshape(self.compressed_weight_shape) result = decompress_symmetric(result, self.scale) result = result.reshape(self.result_shape) if self.result_shape is not None else result - x.weight = result.type(self.result_dtype) + result = result.type(self.result_dtype) + if return_decompressed_only: + return result + else: + x.weight = result