File size: 4,828 Bytes
01906b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 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) |