一、引言
关于 LoRA 的具体理论原理可以参考：图解大模型微调系列之：大模型低秩适配器 LoRA（原理篇）
关于 LoRA 的源码解读实操可以参考：图解大模型微调系列之：大模型低秩适配器 LoRA（源码解读与实操篇）
以上两篇博客写的非常好，强烈建议仔细阅读，本篇博客是对上面第二篇博客中的 LoRA 微软官方源码的进一步解读。包含对源码框架的解读，对源码中训练过程和推理过程的分析，以及对源码一句一句的汉语注释。
重新整理后带汉语注释的 LoRA 源码 Git 链接。
二、LoRA 模型搭建
关于如何通过 LoRA 原理搭建 small GPT2 模型的代码主要分布在 layers.py 和 model.py 两个文件中，其中相关类的调用和继承关系如下图所示，utils.py 中是与模型应用相关的两个函数。
Model/GPT2LMModel()：small GPT2 的完整定义类，包括完整结构定义、loss 计算和上载权重功能。
Model/GPT2Config()：small GPT2 的配置类，仅仅包含众多配置参数。
Model/GPT2LMHead()：small GPT2 的 head，仅仅包含一个反 embedding 层，且这个 embedding 层需要与模型输入端的 emdedding 进行参数绑定。
Model/GPT2Model()：small GPT2 的 backbone，是由位置嵌入层、特征嵌入层和一系列 transformer decoder 层搭建的模型，位置嵌入层负责进行位置编码，特征嵌入层负责将输入的整形序列的中每个整型值转换为词向量特征，所以模型的输入必须是完成切词和映射后的整形序列，例如 [123, 126, 345, ...]。经过嵌入层后送入 decoder 的特征尺寸为 in_features = n_heads*size_embding。
Model/Block()：transformer decoder 模块，包含自注意力机制层、线性映射层和两个层归一化层，是标准的 decoder 结构。
Model/Attention()：transformer decoder 中的自注意力模块，transformer 的核心类，整个注意力机制包括两部分。第一部分是 q k v 的计算，第二部分是注意力结果计算。第一部分融合进了 LoRA 训练策略，可以实现 LoRA 训练，可以合并参数也可以拆分参数。第二部分不在乎有没有 LoRA，就是纯注意力计算的结果。且这里的实现与标准 transformer 不太一样，标准 transformer 在计算 q k v 时每个 head 独自计算，相互不干扰，例如 head1_q=head1_Q*head1_x，不需要 head2_x 参与，而该类在计算时不同 head 相互干扰，具体说就是 head1_q 在计算时不仅仅通过 head1_x 计算，还需要 head2_x，... 等参与计算，例如 [head1_q, head2_q]=Q*[head1_x, head2_x]，这会造成 Q 参数量和计算量的二次幂级增加，原本不同 head 独立计算时，只需要 n_head 个矩阵即可完成不同 head 的 q 的计算，但该类则需要 n_head*n_head 个矩阵组成的大 Q 进行计算。decoder 中的前向注意力机制是通过一个下三角矩阵实现的，只是将后向的注意力部分乘以了系数 0，所以相比 encoder 并没有减少计算量。在计算注意力系数时，所有的 q 和 k 都会参与计算，不会在意输入序列中的实际有效长度，即 padding 部分也会完整的参与到注意力结果的计算中，所以如果想要减少这部分的计算量，需要尽量少的加入 padding 字符。在其 forward 函数中存在一个参数 layer_past，layer_past 在模型训练时为空，在模型推理时，负责缓存前面已完成计算的 q k v。
Layers/MergedLinear()：LoRA 的核心类，输入是 x 输出是 q k v，包含 LoRA 中权重参数 A 和 B 的定义。可以实现 A B 与原始 Weight 的合并和拆分。LoRA 的权重参数 A 和 B 的定义和计算过程也比较特殊，A 和 B 的合并计算是通过分组 conv1D 完成的，仔细分析计算过程，发现其实这个等效于矩阵乘法，其中分组的目的是让 Q K V 的计算相互独立，但不同 head 的计算不独立。A 与 B 的乘法比较特殊是一种交叉乘法，相比原来的不同 head 独立计算，这二次幂级增加了计算量和输出矩阵尺寸(与二次幂增加参数量后的 W 相匹配)，但却不增加 A 和 B 的参数量。所谓交叉计算指的是不同 head 的 A 和 B 的计算不独立，会存在 head_i_A * head_j_B 的情况，这也是计算量和输出参数量增加的原因。该类继承自 nn.Linear，所以只能替换原模型中的 nn.Linear 层。
Layers/LoRALayer()：LoRA 的配置类，只是定义了所需参数，可以配置原模型中哪些层需要加 LoRA Adapter。
Model/MLP()：transformer decoder 中的线性映射层，包含两个线性层(先升维再降维)和一个激活函数。
这里仅仅展示 LoRA 核心类 MergedLinear() 的源码：
# 该类为 LoRA 的核心类，forward 函数输入是某一个注意力模块中的 x，输出是对应 q k v 组成的矩阵。类中定义了原注意力参数 W 和对应的 LoRA 权重 A 和 B。当 merge_weights=True 时，可以通过 train() 函数实现总体参数的合并和拆分，当其等于 True 时，参数默认被拆分。
# 原注意力模块参数 W 在定义后被冻结，不能被训练，train() 函数不影响它是否可训练，该类不修改 W。W 定义在 nn.Linear 中，注意 W 的尺寸为 (in_features, out_features)，其中 in_features=n_heads*size_embding，out_features=3*n_heads*size_embding。
# 这意味着原注意模块中不同 head 的计算是不独立的，这与原 Transformer 不太一样，原来是先不同 head 独立计算，完成计算后再通过全连接层合并多头特征，而这里直接第一步就是非独立计算，这样会极大地增大参数量和计算量，它们随 head 数量二次幂增加，原本是线性增加。
# LoRA 的权重参数 A 和 B 的定义和计算过程也比较特殊，A 和 B 的合并计算是通过分组 conv1D 计算的，仔细分析计算过程，发现其实这个等效于矩阵乘法，其中分组的目的是让 Q K V 的计算相互独立，但不同 head 的计算是不独立的。A 与 B 的乘法比较特殊是一种交叉乘法，
# 相比原来的不同 head 独立计算，这二次幂级增加了计算量和输出矩阵尺寸(与二次幂增加参数量后的 W 相匹配)，但却不增加 A 和 B 的参数量。所谓交叉计算指的是不同 head 的 A 和 B 的计算不独立，会存在 head_i_A * head_j_B 的情况，这也是计算量和输出参数量增加的原因。
class MergedLinear(nn.Linear, LoRALayer):
# LoRA implemented in a dense layer
def __init__(
self,
in_features: int, # 原注意力模块 Q/K/V 矩阵的行/列尺寸(行列尺寸相等，Q K V 尺寸相等) * n_heads， [head1.., head2.., head3..]
out_features: int, # 原注意力模块 Q K V 的输出尺寸之和(等于 3 倍的独自尺寸) * n_heads，先分 Q K V，再分 heads [Q_head1.., Q_head2.., Q_head3.., K_head1.., K_head2.., K_head3.., V_head1.., V_head2.., Q_head3..]
r: int = 0, # LoRA 的秩 # 从上面 in_features 和 out_features 的展示形式很容易判断，输入不区分 Q K V 但输出区分 Q K V，输入和输出都要区分不同 head
lora_alpha: int = 1, # LoRA 的 alpha # 所以从 in_features 映射到 out_features 的矩阵 W 是 Q K V 独立，但不同 head 不独立的
lora_dropout: float = 0., # LoRA 的 Dropout，在 LoRA 分支上是，先 Dropout 再 A 和 B
enable_lora: List[bool] = [False], # 长度为 3，决定 Q K V 是否使用 LoRA，一般定义为 [True, False, True]，即 K 不使用 LoRA
fan_in_fan_out: bool = False, # 决定是否部分矩阵转置一下，设置为 True
merge_weights: bool = True, # 为 True 时，训练时分开参数，测试时合并参数；为 False 时，训练时和测试时均分开参数；决定是训练还是测试由 train 和 eval 模式决定
**kwargs
):
nn.Linear.__init__(self, in_features, out_features, **kwargs) # 调用 nn.Linear 的基函数，实现对原有参数的定义，注意 in_features=feature*n_heads out_features=3*feature*n_heads
LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights) # 调用 LoRA 基函数，实现相关参数初始化输入
assert out_features % len(enable_lora) == 0, 'The length of enable_lora must divide out_features' # out_features=3*feature*n_heads，Q K V 放在一起以及所有 head 放在一起了，len(enable_lora)=3 所以必须要能整除
self.enable_lora = enable_lora # 长度为 3，决定 Q K V 是否使用 LoRA，一般定义为 [True, False, True]，即 K 不使用 LoRA
self.fan_in_fan_out = fan_in_fan_out # 决定是否部分矩阵转置一下，设置为 True
# Actual trainable parameters # ！！！【状态 1】 ！！！ # ！！！【状态 2】 ！！！
if r > 0 and any(enable_lora): # enable_lora 包含 3 True # enable_lora 包含 2 True
self.lora_A = nn.Parameter(
self.weight.new_zeros((r * sum(enable_lora), in_features)) # (3*r, in_features) # (2*r, in_features) 对于 LoRA A 是 Q K V 独自享有自己的部分，计算时相互独立
)
self.lora_B = nn.Parameter(
self.weight.new_zeros((out_features // len(enable_lora) * sum(enable_lora), r)) # (out_features, r) # (2/3*out_features, r) 对于 LoRA B 是 Q K V 独自享有自己的部分，计算时相互独立
)
self.scaling = self.lora_alpha / self.r
# Freezing the pre-trained weight matrix
self.weight.requires_grad = False
# Compute the indices
self.lora_ind = self.weight.new_zeros(
(out_features, ), dtype=torch.bool
).view(len(enable_lora), -1) # (3, out_features/3) # (3, out_features/3) 全 0 矩阵
self.lora_ind[enable_lora, :] = True # 所有元素全 True # 对应行所有元素全 True
self.lora_ind = self.lora_ind.view(-1) # (out_features) # (out_features)
self.reset_parameters()
if fan_in_fan_out:
self.weight.data = self.weight.data.transpose(0, 1) # (in_features, out_features) # (in_features, out_features)
# 参数初始化，尤其注意所有 LoRA 的初始化为全 0
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
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
nn.init.zeros_(self.lora_B)
# 当 Q K V 均添加 LoRA 时，该函数无意义，否则，该函数负责 padding 增加没有 LoRA 的部分
def zero_pad(self, x): # (out_features, in_features) # (2/3*out_features, in_features)
result = x.new_zeros((len(self.lora_ind), *x.shape[1:])) # (out_features, in_features) # (out_features, in_features)
result[self.lora_ind] = x # (out_features, in_features) # (out_features, in_features) 对应部分有值，其他部分为 0
return result
# LoRA 权重参数 A 和 B 的合并是通过分组 conv1D 计算的，仔细分析计算过程，发现其实这个等效于矩阵乘法，其中分组的目的是让 Q K V 的计算相互独立，但不同 head 的计算不独立。与原本不同 head 独立计算相比，A 和 B 的参数量并没有增加，
def merge_AB(self): # 但 A*B 的计算量和输出矩阵尺寸却二次幂级增加了，原因是 A 与 B 的乘法比较特殊是一种交叉乘法，会存在 A_head_i * B_head_j 的情况，并非仅仅包含 A_head_i * B_head_i 和 A_head_j * B_head_j
def T(w):
return w.transpose(0, 1) if self.fan_in_fan_out else w
delta_w = F.conv1d( # F.conv1d # enable_lora 包含 3 True # enable_lora 包含 2 True
self.lora_A.unsqueeze(0), # input： (Batch_Size, In_Channel, Length) # (1, 3*r, in_features) # (1, 2*r, in_features)
self.lora_B.unsqueeze(-1), # weight：(Out_Channel, In_Channel/Group, Kernel_Size) # (out_features, r, 1) # (2/3*out_features, r, 1)
groups=sum(self.enable_lora) # group： Group # 3 # 2
).squeeze(0) # output：(1, Batch_Size, Out_Channel, In_Channel) # (1, out_features, in_features) # (1, 2/3*out_features, in_features) 无 squeeze()
return T(self.zero_pad(delta_w)) # (in_features, out_features) # (in_features, out_features)
# 简单点总结，这个函数的作用是：merge_weights 为 False 时不做任何处理，此时参数一定是拆分状态(前提是 merge_weights 在完成定义后不再被修改)，因为整个代码中，只有这个函数会合并参数，如果这个函数没有去合并参数，那参数一定处于拆分状态
# merge_weights 为 True 时，如果 train mode 为 True，则使得参数处于拆分状态，如果 mode 为 False，则使得参数处于合并状态，所谓参数拆分是指将总 W 拆分为 Pretrain_W 和 LoRA_W
def train(self, mode: bool = True): # mode 决定是训练模式还是测试默认，默认是 True，一般都是使用默认模式，但当设置 false 时，则等效于调用 eval() 函数，此时一般直接调用 eval() 函数
# 没用到的函数
def T(w):
return w.transpose(0, 1) if self.fan_in_fan_out else w
# 让原有参数变为 train 模式，我理解这里其实没有造成任何影响，train 只对 BN、 Dropout 和 LoRA 产生影响，对 LoRA 产生影响时还需要结合 merge_weights 状态的设置，设置为 True 时才产生影响
nn.Linear.train(self, mode)
# 如果是训练模式，这里的模式不影响参数是否可训练，只对 BN、 Dropout 和 LoRA 产生影响，对 LoRA 产生影响时还需要结合 merge_weights 状态的设置，设置为 True 时才产生影响
if mode:
# 如果 merge_weights 为 Ture (意味着训练拆分|测试合并)，当前是训练模式，所以如果当前是合并状态则需要拆分参数，如果当前是拆分状态不需要做任何处理；如果 merge_weights 为 False，则无需做任何处理，因为默认就是拆分状态
if self.merge_weights and self.merged:
# Make sure that the weights are not merged
if self.r > 0 and any(self.enable_lora):
self.weight.data -= self.merge_AB() * self.scaling # (in_features, out_features) # (in_features, out_features)
self.merged = False
# 如果是测试模式，这里的模式不影响参数是否可训练，只对 BN、 Dropout 和 LoRA 产生影响，对 LoRA 产生影响时还需要结合 merge_weights 状态的设置，设置为 True 时才产生影响
else:
if self.merge_weights and not self.merged:
# Merge the weights and mark it
if self.r > 0 and any(self.enable_lora):
self.weight.data += self.merge_AB() * self.scaling # (in_features, out_features) # (in_features, out_features)
self.merged = True
# 执行函数
def forward(self, x: torch.Tensor):
def T(w):
return w.transpose(0, 1) if self.fan_in_fan_out else w
# self.merged 描述的是参数当前合并状态
if self.merged: # 如果参数已经合并，那直接单次计算即可完成输出
return F.linear(x, T(self.weight), bias=self.bias) # (batch, seq_length, out_features) = (batch, seq_length, in_features) * (out_features, in_features)^T
else: # 如果参数还没有合并，那需要分别经过 W 和 delta_W 的计算，然后合并结果
result = F.linear(x, T(self.weight), bias=self.bias) # (batch, seq_length, out_features) = (batch, seq_length, in_features) * (out_features, in_features)^T
if self.r > 0:
result += self.lora_dropout(x) @ T(self.merge_AB().T) * self.scaling # (batch, seq_length, out_features) = (batch, seq_length, out_features) + (batch, seq_length, in_features) * (in_features, out_features)
return result三、数据集
这里并不是从零起训练模型，而是采用 SFT 进行模型微调，如下所示数据集中包含输入和人工标注的输出部分，“||” 分割两部分。
name : Blue Spice | Type : coffee shop | area : city centre||A coffee shop in the city centre area called Blue Spice .
name : Blue Spice | Type : coffee shop | area : city centre||Blue Spice is a coffee shop in city centre .
name : Blue Spice | Type : coffee shop | area : riverside||There is a coffee shop Blue Spice in the riverside area .
name : Blue Spice | Type : coffee shop | area : riverside||At the riverside , there is a coffee shop called The Blue Spice .
name : Blue Spice | Type : coffee shop | customer rating : 5 out of 5 | near : Crowne Plaza Hotel||The coffee shop Blue Spice is based near Crowne Plaza Hotel and has a high customer rating of 5 out of 5 .
Data_utils/FT_Dataset() 负责数据读取，数据需要事先经过分词和映射为整形序列，并在输入数据 conditions 和输出数据 completion 后面都添加一个序列 id 为 50256 的特殊字符，所以送入 Data_utils/FT_Dataset() 的数据如下所示。
{"context": [3672, 1058, 4518, 43537, 930, 5994, 1058, 6891, 6128, 930, 1989, 1058, 1748, 7372, 50256], "completion": [317, 6891, 6128, 287, 262, 1748, 7372, 1989, 1444, 4518, 43537, 764, 50256]}
{"context": [3672, 1058, 4518, 43537, 930, 5994, 1058, 6891, 6128, 930, 1989, 1058, 1748, 7372, 50256], "completion": [4518, 43537, 318, 257, 6891, 6128, 287, 1748, 7372, 764, 50256]}
{"context": [3672, 1058, 4518, 43537, 930, 5994, 1058, 6891, 6128, 930, 1989, 1058, 18180, 485, 50256], "completion": [1318, 318, 257, 6891, 6128, 4518, 43537, 287, 262, 18180, 485, 1989, 764, 50256]}
{"context": [3672, 1058, 4518, 43537, 930, 5994, 1058, 6891, 6128, 930, 1989, 1058, 18180, 485, 50256], "completion": [1629, 262, 18180, 485, 837, 612, 318, 257, 6891, 6128, 1444, 383, 4518, 43537, 764, 50256]}
{"context": [3672, 1058, 4518, 43537, 930, 5994, 1058, 6891, 6128, 930, 6491, 7955, 1058, 642, 503, 286, 642, 930, 1474, 1058, 12223, 68, 23280, 12696, 50256], "completion": [383, 6891, 6128, 4518, 43537, 318, 1912, 1474, 12223, 68, 23280, 12696, 290, 468, 257, 1029, 6491, 7955, 286, 642, 503, 286, 642, 764, 50256]}Data_utils/FT_Dataset() 的输出数据格式是一个样本对应一个字典，字典的键包括 "query"，"query_len"，"input"，"target"，"mask"。模型训练时使用 "input"，"target"，"mask"，模型推理时使用 "query"，"query_len"。"input" 是 conditions+completion+补长字符，不包含其他特殊字符，总长度是 max_seq_length，"target" 是 "input" 向前错一位，"mask" 是 0|1 的序列，用于标识 "input" 中的有效长度。"query" 是 conditions+补长字符，不包含其他特殊字符，总长度是 max_seq_length，"query_len" 是一个整形数值表示 "query" 中的有效长度。所以不管是模型训练时还是模型推理时，送入模型的整形序列长度都是固定的 max_seq_length，与实际有效数据长度无关。
四、LoRA 模型训练
gpt2_ft.py 负责 LoRA 模型训练，其会调用 utils.py，optimizer.py，gpu.py 中的部分函数。训练部分的核心代码在 gpt2_ft/train_validate() 中，其训练过程与标准的 torch 模型训练过程无异，只是采用了分布式训练。
LoRA 模型训练使用的输入是 [conditions+completion+补长字符]，目标输出是错一位的 [conditions+completion+补长字符]，在计算 loss 时基于 "mask" 仅仅计算有效长度 conditions+completion 部分的 loss，不会计算 padding 部分的 loss。模型训练时的输入数据 batch 中不要求所有样本的有效长度都相等，但所有样本最终都会被 padding 到固定长度 max_seq_length，这会存在大量的算力浪费，因为在模型中并非按照有效长度计算注意力结果，而是按照输入数据长度计算注意力结果。在 HuggingFace 的 dataset 库中就实现了一些功能可以节省算力和内存，这些功能是：1. 在组成输入数据 batch 时尽量挑选长度相差不大的样本；2. 允许挑选长度过短的样本组成一个样本进行训练；3. batch 中的样本在 padding 时，不是 padding 到固定长度，而是 padding 到本 batch 中最长样本的长度。当然前提是这里的 small GPT2 和 HuggingFace 中的模型都是不要求输入数据必须满足固定长度的。
五、LoRA 模型推理
gpt2_beam.py 负责 LoRA 模型推理，其会调用 gpu.py 中的部分函数。推理部分的核心代码在 gpt2_beam/beam() 中。
LoRA 模型推理时，首先一次性输入长度为 max_seq_length 的输入序列 [conditions+补长字符]，一个 batch 中不要求所有 conditions 都等长，并有个参数 len 记录 batch 中每个样本的真实长度。模型推理过程中，与训练过程类似，存在大量算力浪费。在第一次推理结束后，模型会输出一个 layer_past 用于缓存本轮计算所得到的 k v。从第二轮推理开始，每个样本依据 len 记录的有效长度从模型上一轮最后一个有效输出端处拿到一个输出字符(作为上一轮推理的输出)，然后将这个字符(长度固定为 1)再送入模型进行推理得到新一个输出字符，这里依然存在小量算力浪费，计算注意力系数时，依然需要与整个 max_seq_length 做计算，而不是仅仅与有效长度做计算。从第二轮推理开始，原已经计算过 k v 的字符不用再次计算，只需从缓存中提取即可，每次仅计算单个新字符的 q k v，计算完成后需要将新 k v 缓存以备下一个新字符的查询，缓存 k v 的数据结构 layer_past 的长度在模型针对这个样本的第一次推理时就固定为输入序列的最大长度了，后续更新缓存时，只需要在上一步真实长度 +1 处填上新的 k v 即可。
LoRA 模型推理时采用的是束搜索方法。具体来说，模型第一次推理时，输入序列尺寸为 [batch*n_beams, max_seq_length]，其中包含 batch 个不同的样本组，一个样本组中的 n_beams 个样本完全相同，是通过复制得来，这里存在大量算力浪费。输出序列尺寸为 [batch, n_beams*vocab]，依然是 batch 个样本组，从每个样本组中挑选 n_beams 个得分最高的样本(得分通过预测的每个新字符的概率累加得到)，得到最终输出序列尺寸为 [batch, n_beams]。从模型第二次推理开始，输入序列尺寸为 [batch*n_beams, 1]，并搭配 layer_past 一起送入模型进行推理。输出序列尺寸为 [batch, n_beams*vocab]，依然是 batch 个样本组，从每个样本组中挑选 n_beams 个得分最高的样本(得分通过预测的每个新字符的概率累加得到)，得到最终输出序列尺寸为 [batch, n_beams]，同时需要根据挑选结果更新 layer_past 中跟踪的样本，即重新组建 beams。循环往复直至达到终止条件，在循环过程中，有可能某个 beam 输出了终止字符 [eos]，那可以把这个 beam 对应的 sentence 作为对应样本的一个输出结果，同时将该 beam 的得分设置为 -inf，这样下次迭代更新 beams 时就会把这个 beam 删除，继续跟踪新的 beam。所以理想情况下，完成所有迭代后，每一个样本都会得到 n 个不同的输出结果，每个结果也会搭配一个得分，这里的 n 一般大于等于 n_beams，达到终止条件后，即使没有输出 [eod] 字符，也认为该样本生成结束。
在每个样本的束搜索过程中，按照得分挑选新一轮 beams 时，提前设置了 3 种后处理条件，以增加生成结果的多样性和真实性。第一，对前面已经被预测过的字符，在本轮的得分上乘以小于 1 的系数，进行惩罚；第二，对于紧挨着前面预测的字符如果与本轮所预测的字符组合起来的新字符组，如果在前面已经被预测过了，则可以对对应字符的得分进行惩罚；第三，如果当前预测句子太短，可以对 [eos] 字符对应的得分进行惩罚。
在模型推理过程中，除了按照 HuggingFace 的方式修改样本 padding 方式可以节省计算量外，还可以将 layer_past 的长度修改可变长度，让长度等于有效长度，避免当前字符的 q 在与所有字符的 k 计算注意力系数时，浪费算力。
