增加lora载入时去除部分模块的正则表达式写法

app.py

@@ -17,7 +17,7 @@ parser.add_argument('--ckpt',type=str,default="rwkv-loramerge-0426-v2-4096-epoch
 parser.add_argument('--model_path',type=str,default=None,help="local model path") 
 parser.add_argument('--lora', type=str, default=None, help='lora checkpoint path')
 parser.add_argument('--lora_alpha', type=float, default=0, help='lora alpha')
-parser.add_argument('--lora_layer_filter',type=str,default=None,help='layer filter. Default merge all layer. Example: "25-31"')
+parser.add_argument('--lora_layer_filter',type=str,default=None,help='layer filter. Default merge all layer. Example: "0.2*25-31"')
 args = parser.parse_args()
 os.environ["RWKV_JIT_ON"] = '1'
 

rwkv_lora.py

@@ -7,11 +7,21 @@ import types, gc, os, time, re
 import torch
 from torch.nn import functional as F
 
+# valid_filter_pattern = r"(((\d+\.\d+\*)?(\d+)(-\d+)?(/\S+)?|(/\S+))(\s+|$))+"
 def get_filter_keys_and_merge_coef(layer_filter):
     if layer_filter:
         layers = []
         layer_coef = {}
+        layer_remove_patterns = {}
         for layer in layer_filter.split(' '):
+            if '/' in layer: #过滤pattern，需要写成正则表达式
+                layer,_,remove_pattern = layer.partition('/')
+                remove_pattern = re.compile(remove_pattern)
+            else:
+                remove_pattern = None
+            if layer=='':
+                layer_remove_patterns['global']=remove_pattern
+                continue
             if '*' in layer:
                 coef,_,layer = layer.partition('*')
                 coef = float(coef)
@@ -20,22 +30,31 @@ def get_filter_keys_and_merge_coef(layer_filter):
             if layer.isdecimal():
                 layers.append(int(layer))
                 layer_coef[int(layer)]=coef
+                layer_remove_patterns[int(layer)]=remove_pattern
             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
+                    layer_remove_patterns[l]=remove_pattern
             else:
                 raise NotImplementedError("layer_filter Not implemented:",layer_filter)
         layers = sorted(set(layers))
-        layer_prefixes = tuple(f"blocks.{l}." for l in layers)
+        # layer_prefixes = tuple(f"blocks.{l}." for l in layers)
         def filter_keys(keys): 
             new_keys = []
             for key in keys:
+                if layer_remove_patterns.get("global") and layer_remove_patterns['global'].search(key):
+                    continue #符合全局去除规则
                 if key.startswith("blocks."): #过滤掉blocks开头，且不在允许范围内的权重
-                    if not key.startswith(layer_prefixes):
+                    l = int(key.split('.')[1])
+                    if l not in layers: #不在允许层，过滤掉
+                        continue
+                    if layer_remove_patterns[l] and layer_remove_patterns[l].search(key): #符合对应层的去除规则，过滤掉
                         continue
+                    # if not key.startswith(layer_prefixes):
+                    #     continue
                 new_keys.append(key)
             return new_keys
         def merge_coef(key):
@@ -59,6 +78,8 @@ def lora_merge(base_model,lora,lora_alpha,device="cuda",layer_filter=None,):
     w_lora: Dict[str, torch.Tensor] = torch.load(lora, map_location='cpu')
     # pdb.set_trace() #DEBUG
     for k in filter_keys(w_lora.keys()): #处理time_mixing之类的融合
+        if k in w:
+            print(f"replacing {k}")
         w[k] = w_lora[k]
     output_w: typing.OrderedDict[str, torch.Tensor] = OrderedDict()
     # merge LoRA weights