class LoRALayer():
def __init__(
self,
r: int,
lora_alpha: int,
lora_dropout: float,
merge_weights: bool,
):
# -----------------------------------------------------
# lora的秩
# -----------------------------------------------------
self.r = r
# -----------------------------------------------------
# 用于计算缩放因子scaling = lora_alpha/r
# -----------------------------------------------------
self.lora_alpha = lora_alpha
# -----------------------------------------------------
# Optional dropout
# -----------------------------------------------------
if lora_dropout > 0.:
self.lora_dropout = nn.Dropout(p=lora_dropout)
else:
self.lora_dropout = lambda x: x
# -----------------------------------------------------
# 表示当前pretrain部分(self.weight)中是否已经融入了lora部分
# self.merged=True,pretrain部分已包含lora,
# 则进行forward是可以直接用pretrain部分
# self.merged=False, pretrain部分未包含lora,
# 则进行forward时需要用pretrain+lora的结果
# 【表示合不合这个状态】
# -----------------------------------------------------
self.merged = False
# -----------------------------------------------------
# self.merge_weights = True,遵从训练时拆分pretain与lora,
# 推理时合并pretrain与lora
# self.merge_weights = False, 遵从无论是训练还是推理,
# 都拆分pretrain与lora
# 【表示合不合这个意愿】
# -----------------------------------------------------
self.merge_weights = merge_weights
class MergedLinear(nn.Linear, LoRALayer):
# LoRA implemented in a dense layer
def __init__(
self,
in_features: int,
out_features: int,
r: int = 0,
lora_alpha: int = 1,
lora_dropout: float = 0.,
enable_lora: List[bool] = [False],
fan_in_fan_out: bool = False,
merge_weights: bool = True,
**kwargs
):
# --------------------------------------------------------------------
# in_features: hidden_size
# out_features: hidden_size*3
# 将nn.linear和LoRALayer的属性继承过来
# --------------------------------------------------------------------
nn.Linear.__init__(self, in_features, out_features, **kwargs)
LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
merge_weights=merge_weights)
# --------------------------------------------------------------------
# 因为我们是把QKV三个矩阵的计算合成一个来计算的,
# 因此out_features要能被矩阵个数整除
# --------------------------------------------------------------------
assert out_features % len(enable_lora) == 0, \
'The length of enable_lora must divide out_features'
# --------------------------------------------------------------------
# 表示哪一块矩阵是需要做lora的,例如[True, False, True],
# 说明第一和第三块需要,第二块不要
# --------------------------------------------------------------------
self.enable_lora = enable_lora
# --------------------------------------------------------------------
# bool值,表示是否需要对矩阵做转置,下文中会看到细节
# --------------------------------------------------------------------
self.fan_in_fan_out = fan_in_fan_out
# --------------------------------------------------------------------
# Actual trainable parameters
# 如果秩>0,且其中有任意一矩阵需要做lora
# --------------------------------------------------------------------
if r > 0 and any(enable_lora):
# --------------------------------------------------------------------
# 初始化A矩阵
# self.weight继承自nn.linear下的属性,shape=(in_features, out_features)
# =(hidden_size, hidden_size*3)
# lora_A的shape=(r*需要做lora的矩阵数量, hidden_size)
# 使用new_zeros,可以复制原来矩阵的数据类型和存储设备
# --------------------------------------------------------------------
self.lora_A = nn.Parameter(
self.weight.new_zeros((r * sum(enable_lora), in_features)))
# --------------------------------------------------------------------
# 初始化B矩阵
# lora_B的shape=(hidden_size*需要做lora的矩阵数量, r)
# --------------------------------------------------------------------
self.lora_B = nn.Parameter(
self.weight.new_zeros((out_features // len(enable_lora) * sum(enable_lora), r))
) # weights for Conv1D with groups=sum(enable_lora)
# --------------------------------------------------------------------
# lora权重的缩放因子,
# 一方面决定预训练知识和lora学得的新知识间的比例,另一方面通过1/r因子使得
# 训练过程更加稳定(参考本文3.2部分的说明,以及lora原理篇中的讲解)
# 在gpt2做nlg的任务里,作者将lora_alpha设置为32,r设置为4
# --------------------------------------------------------------------
self.scaling = self.lora_alpha / self.r
# --------------------------------------------------------------------
# Freezing the pre-trained weight matrix
# 原始的QKV矩阵需要冻结
# --------------------------------------------------------------------
self.weight.requires_grad = False
# --------------------------------------------------------------------
# Compute the indices
# 1. lora_ind的shape=(3, hidden_size),默认值为False
# 2. 对lora_ind,只有需要做lora矩阵的那些行,才填充为True,例如hidden_size = 5,
# 第一和第三个矩阵需要做lora,则lora_ind变为:
# tensor([[ True, True, True, True, True],
# [False, False, False, False, False],
# [ True, True, True, True, True]])
# 最后,将lora_ind的形式转成shape=(3*hidden_size)
# --------------------------------------------------------------------
self.lora_ind = self.weight.new_zeros(
(out_features, ), dtype=torch.bool
).view(len(enable_lora), -1)
self.lora_ind[enable_lora, :] = True
self.lora_ind = self.lora_ind.view(-1)
# 对lora_A和lora_B重新做参数初始化
self.reset_parameters()
# --------------------------------------------------------------------
# 若fan_in_fan_out = True,
# 则把weight的shape(hidden_size, 3*hidden_size) ->
# (3*hidden_size, hidden_size)
# 这样做的原因是,例如输入x=(b,s,h), w=(h, 3h)
# 调用F.linear(x, w)时,默认执行的是(b,s,h)*(3h,h),
# 因此要通过这个方式,把w先做转置
# --------------------------------------------------------------------
if fan_in_fan_out:
self.weight.data = self.weight.data.transpose(0, 1)
def reset_parameters(self):
"""
重新初始化参数
"""
nn.Linear.reset_parameters(self)
if hasattr(self, 'lora_A'):
# --------------------------------------------------------------------
# initialize A the same way as the default for nn.Linear and B to zero
# 对lora_A,采用kaiming_uniform; 对lora_B,初始化为0
# --------------------------------------------------------------------
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
nn.init.zeros_(self.lora_B)
def zero_pad(self, x):
"""
x是数据过B矩阵后的结果,x.shape =(b, s, hidden_size*需要做lora的矩阵数量)
我们要将这个结果做zero padding,使得:
x.shape = (b, s, hidden_size*总矩阵数量),在QKV的例子中,总矩阵数量=3,
其中不是我们要做lora矩阵的部分,都padding上0
"""
result = x.new_zeros((*x.shape[:-1], self.out_features))
result = result.view(-1, self.out_features)
result[:, self.lora_ind] = x.reshape(
-1, self.out_features // len(self.enable_lora) * sum(self.enable_lora)
)
return result.view((*x.shape[:-1], self.out_features))
def train(self, mode: bool = True):
def T(w):
return w.transpose(0, 1) if self.fan_in_fan_out else w
nn.Linear.train(self, mode)
# --------------------------------------------------------------
# 如果是model.train()模式
# --------------------------------------------------------------
if mode:
# --------------------------------------------------------------
# 参见4.1讲解
# --------------------------------------------------------------
if self.merge_weights and self.merged:
# Make sure that the weights are not merged
if self.r > 0 and any(self.enable_lora):
# ---------------------------------------------------
# delta_w: lora_A * lora_B矩阵后的结果
# 输入x * delta_w是最终lora矩阵的输出
# lora_A: shape=(r*需要做lora的矩阵数, h)
# lora_B: shape=(h*需要做lora的矩阵数, r)
# 过F.conv1d(lora_A当成输入,lora_B当成卷积)后的:
# delta_w = (1, h*需要做lora的矩阵数, h)
# squeeze之后delta_w = (h*需要做lora的矩阵数, h)
# ---------------------------------------------------
delta_w = F.conv1d(
self.lora_A.data.unsqueeze(0),
self.lora_B.data.unsqueeze(-1),
groups=sum(self.enable_lora)
).squeeze(0)
# 从pretrain矩阵中剔除lora部分
self.weight.data -= self.zero_pad(T(delta_w * self.scaling))
# 当原来的lora剥离出来后,self.merge要设置成False,表示此时pretrain中不再有lora
self.merged = False
# --------------------------------------------------------------
# 如果是train.eval()模式: model.train(False) = train.eval()
# --------------------------------------------------------------
else:
# --------------------------------------------------------------
# 参见4.1讲解
# --------------------------------------------------------------
if self.merge_weights and not self.merged:
# Merge the weights and mark it
if self.r > 0 and any(self.enable_lora):
delta_w = F.conv1d(
self.lora_A.data.unsqueeze(0),
self.lora_B.data.unsqueeze(-1),
groups=sum(self.enable_lora)
).squeeze(0)
self.weight.data += self.zero_pad(T(delta_w * self.scaling))
# 既然已将lora合进pretrain,那么就要把self.merged设置为True
self.merged = True
def forward(self, x: torch.Tensor):
"""
Params:
x: 输入数据,形状为(b, s, h)
"""
# --------------------------------------------------------------------
# 定义矩阵专置函数
# 例如原始矩阵是(hidden_size, 3*hidden_size),
# 转置之后为(3*hidden_size, hidden_size)
# --------------------------------------------------------------------
def T(w):
return w.transpose(0, 1) if self.fan_in_fan_out else w
# --------------------------------------------------------------------
# self.merged = True,如果lora和pretrain部分已经合起来
# 返回结果的shape = (b, s, 3*hidden_size)
# --------------------------------------------------------------------
if self.merged:
return F.linear(x, T(self.weight), bias=self.bias)
# --------------------------------------------------------------------
# self.merged = False,如果lora和pretrain部分没有合起来
# --------------------------------------------------------------------
else:
# 先用merged后的矩阵正常计算,result的shape = (b, s, 3*hidden_size)
result = F.linear(x, T(self.weight), bias=self.bias)
if self.r > 0:
# --------------------------------------------------------------------
# 数据过lora_A后的结果
# 结果shape = (b, s, h) * (h, r*需要做lora的矩阵的数量)
# = (b, s, r*需要做lora的矩阵的数量)
# --------------------------------------------------------------------
after_A = F.linear(self.lora_dropout(x), self.lora_A)
# --------------------------------------------------------------------
# 数据过lora_A,再过lora_B后的结果
# F.conv1d(input=输入数据,weight=卷积核)
# input = 输入数据过lora_A后的结果,shape = (b, r*需要做lora矩阵的数量, s)
# weight = (hidden_size*需要做lora的矩阵数量, r, 1)
# groups = 需要做lora的矩阵数量
# 通过F.conv1d计算出after_B的shape = (b, hidden_size*需要做lora的矩阵数量, s)
# 做transpose(-2, -1)转置后after_B的shape = (b, s, hidden_size*需要做lora的矩阵数量)
# --------------------------------------------------------------------
after_B = F.conv1d(
after_A.transpose(-2, -1),
self.lora_B.unsqueeze(-1),
groups=sum(self.enable_lora)
).transpose(-2, -1)
# --------------------------------------------------------------------
# 数据最终结果 = pretrain + lora
# zero_pad:将过B矩阵后的数据从shape=(b, s, hidden_size*需要做lora的矩阵数量)
# 变成(b, s, hidden_size*总矩阵数量),
# 其中不属于需要做lora的部分就用0填重
# self.scaling = alpha/r,但暂时不知道scaling的作用
# --------------------------------------------------------------------
result += self.zero_pad(after_B) * self.scaling
return result
- 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.
- 152.
- 153.
- 154.
- 155.
- 156.
- 157.
- 158.
- 159.
- 160.
- 161.
- 162.
- 163.
- 164.
- 165.
- 166.
- 167.
- 168.
- 169.
- 170.
- 171.
- 172.
- 173.
- 174.
- 175.
- 176.
- 177.
- 178.
- 179.
- 180.
- 181.
- 182.
- 183.
- 184.
- 185.
- 186.
- 187.
- 188.
- 189.
- 190.
- 191.
- 192.
- 193.
- 194.
- 195.
- 196.
- 197.
- 198.
- 199.
- 200.
- 201.
- 202.
- 203.
- 204.
- 205.
- 206.
- 207.
- 208.
- 209.
- 210.
- 211.
- 212.
- 213.
- 214.
- 215.
- 216.
- 217.
- 218.
- 219.
- 220.
- 221.
- 222.
- 223.
- 224.
- 225.
- 226.
- 227.
- 228.
- 229.
- 230.
- 231.
- 232.
- 233.
- 234.
- 235.
- 236.
- 237.
- 238.
- 239.
- 240.
- 241.
- 242.
- 243.
- 244.
- 245.
- 246.
- 247.
- 248.
- 249.
- 250.
- 251.
- 252.
- 253.
- 254.
- 255.
- 256.
- 257.
- 258.
- 259.
- 260.
- 261.
- 262.
- 263.
- 264.
- 265.
- 266.
- 267.
- 268.
- 269.
- 270.
- 271.
- 272.
- 273.
- 274.
- 275.
- 276.
- 277.
- 278.
- 279.
- 280.
- 281.
- 282.
- 283.
- 284.