lora filter 允许比例，例如0.2*0-4

rwkv_lora.py

@@ -7,38 +7,53 @@ import types, gc, os, time, re
 import torch
 from torch.nn import functional as F
 
-def get_filter_keys(layer_filter):
+def get_filter_keys_and_merge_coef(layer_filter):
     if layer_filter:
         layers = []
+        layer_coef = {}
         for layer in layer_filter.split(' '):
+            if '*' in layer:
+                coef,_,layer = layer.partition('*')
+                coef = float(coef)
+            else:
+                coef = 1 
             if layer.isdecimal():
                 layers.append(int(layer))
+                layer_coef[int(layer)]=coef
             elif '-' in layer:
                 start,_,end = layer.partition('-')
                 start,end = int(start),int(end)
                 layers.extend(range(start,end+1))
+                for l in range(start,end+1):
+                    layer_coef[l] = coef
             else:
                 raise NotImplementedError("layer_filter Not implemented:",layer_filter)
         layers = sorted(set(layers))
         layer_prefixes = tuple(f"blocks.{l}." for l in layers)
-        def filter_keys(keys):
+        def filter_keys(keys): 
             new_keys = []
             for key in keys:
-                if key.startswith("blocks."):
+                if key.startswith("blocks."): #过滤掉blocks开头，且不在允许范围内的权重
                     if not key.startswith(layer_prefixes):
                         continue
                 new_keys.append(key)
             return new_keys
-
+        def merge_coef(key):
+            if key.startswith('blocks.') and int(key.split('.')[1]) in layer_coef:
+                return layer_coef[int(key.split('.')[1])]
+            else:
+                return 1
     else:
         def filter_keys(keys):
             return keys
-    return filter_keys
+        def merge_coef(key):
+            return 1
+    return filter_keys,merge_coef
 
 def lora_merge(base_model,lora,lora_alpha,device="cuda",layer_filter=None,):
     print(f"Loading LoRA: {lora}")
     print(f"LoRA alpha={lora_alpha}, layer_filter={layer_filter}")
-    filter_keys = get_filter_keys(layer_filter)
+    filter_keys,merge_coef = get_filter_keys_and_merge_coef(layer_filter)
     w: Dict[str, torch.Tensor] = torch.load(base_model, map_location='cpu')
     # merge LoRA-only slim checkpoint into the main weights
     w_lora: Dict[str, torch.Tensor] = torch.load(lora, map_location='cpu')
@@ -61,7 +76,7 @@ def lora_merge(base_model,lora,lora_alpha,device="cuda",layer_filter=None,):
                 w[k] = w[k].to(device=device)
                 w[lora_A] = w[lora_A].to(device=device)
                 w[lora_B] = w[lora_B].to(device=device)
-                w[k] += w[lora_B] @ w[lora_A] * (lora_alpha / lora_r)
+                w[k] += w[lora_B] @ w[lora_A] * (lora_alpha / lora_r) * merge_coef(k)
                 output_w[k] = w[k].to(device='cpu', copy=True)
                 del w[k]
                 del w[lora_A]