Upload 2 files

merge.bash

@@ -0,0 +1,7 @@
+python3 my_merge_lora.py\
+    --base_model ~/autodl-tmp/RWKV-4-Raven-7B-v11-Eng49%-Chn49%-Jpn1%-Other1%-20230430-ctx8192.pth\
+    --lora_load ~/RWKV_Train/save_7B_128/rwkv-12.pth\
+    --save ~/autodl-tmp/lora_7B.pth\
+    --lora_schedule poly\
+    --plot\
+    --lora_strategy "0:64,0.3:32,0.4:0,0.5:32,1:64"

my_merge_lora.py

@@ -0,0 +1,151 @@
+from matplotlib import pyplot as plt
+import numpy as np
+from collections import OrderedDict
+import os
+from tqdm import tqdm
+import sys
+from typing import Dict
+import typing
+import torch
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument("--base_model", type=str, required=True)
+parser.add_argument("--lora_load", type=str, required=True)
+parser.add_argument("--lora_schedule", type=str, default="const", choices=["const", "linear", "poly"])
+parser.add_argument("--lora_alpha", type=float, default=None)
+parser.add_argument("--lora_beta", type=float, default=None)
+parser.add_argument("--lora_strategy", type=str, default=None)
+parser.add_argument("--plot", action="store_true")
+parser.add_argument("--save", required=True, type=str)
+args = parser.parse_args()
+
+def f_const(alpha):
+    def inner(x):
+        return alpha
+    return inner
+def f_linear(alpha, beta):
+    def inner(x):
+        return alpha+x*(beta-alpha)
+    return inner
+
+def f_polyline(*args):
+    args = sorted(args)
+    xs = []
+    ys = []
+    xys = args
+    for x, y in args:
+        xs.append(x)
+        ys.append(y)
+    def inner(x):
+        if(x<=xs[0]):
+            return ys[0]
+        if(x>=xs[-1]):
+            return ys[-1]
+        for i, x1 in enumerate(xs):
+            x0 = xs[i-1]
+            if(x0<=x and x<=x1):
+                y0 = ys[i-1]
+                y1 = ys[i]
+                n = (x-x0)/(x1-x0)
+                y = y0*(1-n)+y1*n
+                return y
+        return ys[-1]
+    return inner
+
+
+with torch.no_grad():
+    base = args.base_model
+    lora = args.lora_load
+    w: Dict[str, torch.Tensor] = torch.load(base, map_location="cpu")
+    w_lora: Dict[str, torch.Tensor] = torch.load(lora, map_location="cpu")
+    out: typing.OrderedDict[str, torch.Tensor] = OrderedDict()
+
+    tuned_weights = []
+
+    lora_rank = None
+    listkeys = list(w.keys())
+    for key in listkeys:
+        if key.endswith(".weight"):
+            pref = key[:-len(".weight")]
+            key_A = pref+".lora_A"
+            key_B = pref+".lora_B"
+            if(key_A in w_lora and key_B in w_lora):
+                print("Weight %s will be merged with lora"%pref)
+                tuned_weights.append(key)
+
+                if(lora_rank is None):
+                    lora_A = w_lora[key_A]
+                    lora_B = w_lora[key_B]
+                    lora_rank = lora_A.shape[0]
+                    assert lora_rank == lora_B.shape[1]
+    
+    
+    if(args.lora_schedule == "const"):
+        alpha = args.lora_alpha
+        if(alpha is None):
+            print("assuming lora_alpha = lora_rank = %s"%lora_rank)
+            alpha = lora_rank
+        f_schedule = f_const(alpha)
+    elif(args.lora_schedule == "poly"):
+        xys = []
+        for i in args.lora_strategy.split(","):
+            xy = [float(_) for _ in i.split(":")]
+            xys.append(xy)
+        f_schedule = f_polyline(*xys)
+    else:
+        alpha = args.lora_alpha
+        beta = args.lora_beta
+        if(alpha is None):
+            print("assuming lora_alpha = lora_rank = %s"%lora_rank)
+            alpha = lora_rank
+        if(beta is None):
+            print("assuming beta = lora_rank = %s"%lora_rank)
+            beta = lora_rank
+        f_schedule = f_linear(alpha, beta)
+    
+    weight2alpha = {}
+    for idx, i in enumerate(tuned_weights):
+        # assuming weights are fetched in depth order
+        depth = idx
+        if(len(tuned_weights)>1):
+            depth = depth/(len(tuned_weights)-1)
+        else:
+            depth = 0.5
+        alpha = f_schedule(depth)
+        print("%s: alpha=%.3f, alpha/rank=%.2f"%(i, alpha, alpha/lora_rank))
+        weight2alpha[i] = alpha
+    if(args.plot):
+        tot = 32
+        xs = np.linspace(0, tot-1, tot)
+        ys = [f_schedule(x/(tot-1)) for x in xs]
+        plt.plot(xs, ys)
+        plt.savefig(args.save+".png")
+        print("plot", args.save+".png")
+    for key in tqdm(listkeys):
+        if(key in tuned_weights):
+            pref = key[:-len(".weight")]
+            key_A = pref+".lora_A"
+            key_B = pref+".lora_B"
+            
+            lora_A = w_lora[key_A]
+            lora_B = w_lora[key_B]
+            lora_rank = lora_A.shape[0]
+            assert lora_rank == lora_B.shape[1]
+
+            delta_w = (lora_B @ lora_A)*weight2alpha[key]/lora_rank
+            new_w = w[key]+delta_w
+            out_w = new_w.to(device=w[key].device, dtype=w[key].dtype, copy=True)
+            out[key] = out_w
+            del w[key]
+        else:
+            # print("%s = original"%(key))
+            out[key] = w[key].clone()
+            del w[key]
+    # assert list(out.keys()) == list(w.keys())
+    # for key in out.keys():
+    #     assert out[key].shape == w[key].shape
+    #     assert out[key].dtype == w[key].dtype
+    
+    
+    torch.save(out, args.save)
+    print("saved", args.save)