NNCF lora support

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  

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

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