LLM微调(Finetune)技术--LoRA
作者： AINLP 来源： AINLP
0 前言
1 大模型微调技术原理概述
1.1 Adapter
1.2 P-Tuning
1.3 LST
1.4 LoRA
1.5 小结
2 LoRA代码解析
2.1 MergedLinear源码解析
2.2 对Llama 进行LoRA微调
参考
0 前言
最近因为工作需要，在接触一些大模型微调训练相关的算子实现，因为以往接触inference相关比较多，而对于training相关的技术接触的相对较少，所以本文就以LoRA: Low-Rank Adaptation of Large Language Models 为例来学习一下大模型微调的一些技术细节。
这里依然先给出LoRA的paper 和 code的链接：
-
paper：https://arxiv.org/abs/2106.09685
-
code：https://github.com/microsoft/LoRA
-
对llama使用lora：https://github.com/Lightning-AI/lit-llama/blob/main/finetune/lora.py
另外，也感谢苏神的博客科学空间 解答了我很多困惑～
1 大模型微调技术原理概述
我们知道自ChatGPT爆火以来，国内外科技公司都开始重兵部署在LLM上，比如Meta的Llama、国内的ChatGLM 和Qwen等，这些模型动辄几十B(Billion)的参数量，以70B的模型为例，假设模型参数都以FP16数据类型保存，那么光weight数据就需要70x10^9 * 2 Byte = 130GB 的显存来存放weight的数据，那对于训练而言，所需显存更大，所以如果要对所有weight进行微调训练，这无疑是非常昂贵的。所以，大模型微调技术则是希望只通过训练少量参数就能实现模型的迁移微调，而LoRA就是目前大模型比较主流的一种微调技术。除了LoRA之外还有其他的几种微调技术，本文在正式进入今天的主角LoRA之前，先对其他方法做一个概述。
1.1 Adapter
Parameter-Efficient Transfer Learning for NLP –Adapter方法的思路是用一个较小的神经网络模块插入到模型的不同层, 如下图右侧的Adapter Layer 由两个线性映射层和一个非线性变换层组成，先降维再升维，两个线性映射的可学习参数量比原模型的attention和feed-forward 中的参数少的多。
模型迁移微调的时候，只训练Adapter Layer的相关参数和原始模型中的norm层和最后的分类层。
Adapter
1.2 P-Tuning
通过对LLM的Embedding层进行改造，在微调时固定其他层的weight，只对Embedding层进行重训练 ,常见的Prompt Tuning方法有：Prefix-Tuning：Optimizing Continuous Prompts for Generation、The Power of Scale for Parameter-Efficient Prompt Tuning等
P-Tuning
对于P-Tuning微调的详解可参见详解大模型微调方法Prompt Tuning(内附实现代码)
1.3 LST
LST: Ladder Side-Tuning for Parameter and Memory Efficient Transfer Learning–LST
在简述LST方法之前，我们先思考一下，上述的Adapter方法和P-Tuning方法在训练微调时到底高效在哪？
我们知道在对模型训练时，我们姑且可以笼统的分为两步：反向传播和梯度下降，以Pytorch为例具体就是
loss.backward() #反向传播
optimizer.step() #梯度下降
-
backward() 会根据模型的前向推理计算图来反向的对各个layer中的weight求的偏导用于之后的梯度下降，对各个layer中的input求偏导用于向前层传递梯度，链式求导。
-
step() 则会根据所选用的优化器，对需要训练的参数执行相应的梯度下降策略，我们姑且可以将此过程简单描述成 如下公式 , 其中为学习率
OK，有此基础我们直接进入LST方法的介绍，还是用图说话，LST直接在原模型的推理旁路上加了一个新的分支，也是固定原模型中的参数 ，将原模型各层输出与新建的旁路分支结合 得到输出。
LST
从图中不难发现，三种微调方法：
-
对于Adapter来说虽然只训练新插入的少部分参数，但是整个梯度回传的过程不能省略，换句话说，与微调整个模型相比：1)对于反向传播 过程而言，各层对weight的梯度不用算了，但是对于input的梯度得算(要向前层传递); 2)对于梯度下降 过程而言，只需要下降少部分新插入的层的weight，原模型的weight都不动
-
对于P-Tuning虽然只需要训练Embedding层，但是Embedding层是输入层，所以与微调整个模型相比：1）对于反向传播过程而言，各层对weight的梯度不用算了，但是对于input的梯度得算(要向前层传递); 2)对于梯度下降过程而言，只需要下降Embedding层参数，原模型的weight都不动
所以我们再来回答本小节开始的问题：高效微调的高效 ，我觉得在于减少了所需梯度下降过程 的权重量和计算量，对于反向传播 的过程，不需要保存对原始weight得梯度也就节省了显存 ，但是反向传播的复杂度并没有降低。
- 而对于LST便是一种反向传播过程和梯度下降过程都高效的微调方法，如上图(c)而言，不难发现，LST的反向传播和梯度下降过程都与原始模型无关，相当于我重新定义了一个小的模型结构，通过获取原模型的输出作为输入来协助微调最终的结果
1.4 LoRA
LoRA: Low-Rank Adaptation of Large Language Models – LoRA方法应该是目前针对大语言模型中微调效果最好的一种方法，该方法的示意图如下，具体来说就是固定原始模型权重，然后定义两个低秩矩阵作为新增weight参与运算，并将两条链路的结果求和后作为本层的输出，而在微调时，只梯度下降新增的两个低秩矩阵。
LST
以单个Linear层()举例，用公式表达就是假设原始预训练模型的权重为 , 则定义两个低秩矩阵 和 ，其中 ,并将Linear层的计算过程有原始 调整为
且对矩阵使用随机高斯初始化 ，对使用零初始化 , 那么能保证在微调开始之前。进一步而言，也可以使用 作为缩放参数来调节, 通过调节缩放比例可以调节预训练模型与LoRA的加权占比。
相对于Adapter和LST，LoRA有个好处就是不会引入额外的推理延迟，因为前者相当于在原始模型结构上新增了一些结构(或者说FLOPs) ，而对于LoRA而言的结果的shape与的shape是一致的，也就是说微调完成后是可以将的结果直接累加到原始weight中的，即，这样一来，在推理时微调前后的FLOPs的是一样的
同样，我们再来分析一下微调时反向传播阶段与梯度下降阶段的复杂度，还是以单个Linear层()为例，不带LoRA 时Linear的反向传播时梯度计算公式如下：
带上LoRA 以后Linear的反向传播时梯度计算公式如下：
不难发现：
-
在反向传播阶段 ：同一层，带LoRA相对于不带LoRA所需要求解的梯度计算量还要多一些 ，而因为 的缘故，所以 不是所有可学习的层都带LoRA ，原论文中只对生成的三个Linear带了LoRA，不带LoRA的层只需要求对输入的梯度往回传即可。
-
同理，在梯度下降阶段 ：因为可训练的参数量少，所以需要梯度下降的参数也少，只需要对LoRA的两个低秩矩阵梯度下降即可
所以总的来说，LoRA在反向传播阶段计算复杂度还要多一些，只是需要梯度下降的参数少，所以节省显存，梯度下降的也快。
最后，关于和这两个低秩矩阵的初始化问题，首先我们是需要的结果初始是0，这样就能保证微调开始时新引入的低秩矩阵不会对最终结果造成影响，那么最直接的方式就是令其中一个低秩矩阵初始阶段为全零，另一个为非全零即可，两者不能都为全零，通过带上LoRA的梯度计算公式 (2)如果和这两个低秩矩阵初始化都是全0那么两个矩阵的梯度都是0，也就训不起来。
1.5 小结
微调技术反向传播梯度下降推理延迟
Adapter 不用求原W梯度，得求全层的X梯度 下降少量新增网络的W 会增加
P-Tuning 不用求原W梯度，得求全层的X梯度 下降Embedding层的W 增加较少
LST 只需要求一个小网络的W和X梯度 下降一个新增轻量级网络的W 会增加
LoRA 既要求W梯度，还得求全层的X梯度，且计算量增多 下降少量的W 没有
相对于全局微调，这些高效微调技术为什么会使得训练的速度变快呢？[1] 1. 只更新了部分参数：这就意味着反向传播时不需要保存所有W的梯度(省显存)，此外梯度下降的计算量变少了
-
减少了通信时间：由于更新的参数量变少了，所以（尤其是多卡训练时）要传输的数据量也变少了，从而减少了传输时间；
-
采用了各种低精度加速技术，如FP16、FP8或者INT8量化等。是的你没听错，量化模型也可以高效的迁移微调，如QLoRA
2 LoRA代码解析
关于代码阅读可以直接去看微软的官方LoRA源码https://github.com/microsoft/LoRA，但是我更推荐Lightning-AI的源码https://github.com/Lightning-AI/lit-llama ，它在官方LoRA源码的基础上做了详细的代码注释(实在是详细，这里给个大赞)，也有将LoRA应用在Llama模型上的代码
对于lit-llama 代码我们需要关注这几个文件：
|-- finetune
|-- lora.py #llama 使用lora进行微调的脚本
|-- lit_llama
|-- lora.py #lora方法核心的Class定义文件
|-- model.py #llama 模型定义文件
2.1 MergedLinear源码解析
LoRA方法核心的Class–MergedLinear代码解析，为了节省篇幅我对代码做了些裁剪，这部分代码在lit_llama/lora.py , 完整源码可去github上查看
# 这是带LoRA方法的并且融合了生成QKV三个Linear的Class
class MergedLinear(nn.Linear, LoRALayer):
# LoRA implemented in a dense layer
def __init__(
self,
# ↓ this part is for pretrained weights
in_features: int, # ⚬ in_features: (embeddings_size)
out_features: int, # ⚬ out_features: (3 * embedding_size)
# ↓ the remaining part is for LoRA
r: int = 0, # ⚬ r: 2 可调节
lora_alpha: int = 1, # ΔW缩放因子
lora_dropout: float = 0.,
enable_lora: List[bool] = [False],# 用于确认生成QKV的三个哪个带Lora哪个不带
fan_in_fan_out: bool = False,
merge_weights: bool = True, #是否合并LoRA矩阵与原始预训练模型的权重，在推理阶段最好是将两者合并
**kwargs
):
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)
self.enable_lora = enable_lora
self.fan_in_fan_out = fan_in_fan_out
#lit_llama的代码还给了一些shape的注释，写的很好我就直接用了
# Actual trainable parameters
# To better understand initialization let's imagine that we have such parameters:
# ⚬ in_features: 128 (embeddings_size)
# 注意：这里是为方便理解，假定embeddings_size=128，最终还是要以实际模型为准如llama2-7B这里是4096
# ⚬ out_features: 384 (3 * embedding_size)
# ⚬ r: 2
# ⚬ enable_lora: [True, False, True] #给生成Q和V的weight带Lora，生成K的不带
if r > 0 and any(enable_lora):
self.lora_A = nn.Parameter(
self.weight.new_zeros((r * sum(enable_lora), in_features))) # (4, 128)
self.lora_B = nn.Parameter(
self.weight.new_zeros((out_features // len(enable_lora) * sum(enable_lora), r)) # (256, 2)
) # weights for Conv1D with groups=sum(enable_lora)
# Notes about shapes above
# - self.lora_A has shape (4, 128): 4 because rank is 2 and LoRA is applied only to two matrices;
# 128 is the input size of the x (embedding size). (4, 128) and not (128, 4) because later on in
# 这里要注意weight的转置问题
# F.linear function weights are automatically transposed. In addition conv1d requires channels to be before seq length
# - self.lora_B has shape (256, 2): 256 because LoRA is applied only to two matrices, so the output is 128*2; 2 tells to have two channels per group for group convolution
#ΔW缩放因子
self.scaling = self.lora_alpha / self.r
# Freezing the pre-trained weight matrix
# 此时的self.weight是指继承父类nn.Linear中的weight 也就是预训练模型的weight
self.weight.requires_grad = False # (384, 128)
# Compute the indices
# Indices are needed to properly pad weight updates with zeros. If we want to fine-tune queries and values,
# but not keys, then the weights update should be:
#
# [[ΔW,ΔW,ΔW, ..., 0,0,0, ..., ΔW,ΔW,ΔW,],
# [....................................],
# [ΔW,ΔW,ΔW, ..., 0,0,0, ..., ΔW,ΔW,ΔW,]]
# ↑ ↑ ↑
# ________________________________________
# | query | key | value |
# ----------------------------------------
self.lora_ind = self.weight.new_zeros(
(out_features, ), dtype=torch.bool
).view(len(enable_lora), -1) # (3, 128)
self.lora_ind[enable_lora, :] = True # (3, 128)
self.lora_ind = self.lora_ind.view(-1) # (384,)
def reset_parameters(self):
"""Reset all the weights, even including pretrained ones."""
nn.Linear.reset_parameters(self)
# 初始化
if hasattr(self, 'lora_A'):
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
nn.init.zeros_(self.lora_B)
其中核心的几个方法，我们单拎出来看
def train(self, mode: bool = True):
"""Set the module into train or eval mode if `mode` is True of False respectively.
For train mode (train(True)) if weights are merged we need to subtract weights updates (LoRA_A @ LoRA_B) from pretrained weights so we can continue training LoRA's matrices A and B and keep pretrained weights frozen.
For eval mode (train(False)) if weights are not merged we need to add weight updates to pretrained weights in order to reduce computational overhead during inference.
Args:
mode: if True the module will be set into train mode (affects Dropout and BatchNorm), if False - eval mode.
"""
def T(w):
return w.T if self.fan_in_fan_out else w
# despite being called from nn.Linear this method will put all layers into train mode, including nn.Dropout of course except parameters (such as self.lora_A, self.lora_B)
nn.Linear.train(self, mode)
# 为什么这里要这么写，是因为整个训练过程包括训练和验证，训练阶段LoRA是拆开的
# 验证阶段相当于是推理阶段，所以是要把LoRA与预训练权重合并的
# 训练与验证彼此迭代，所以需要整这么个逻辑
# if train(True) -> unmerge unless we already have them unmerged
# if train(False) -> merge unless we already have them merged
# if train(True)说明是训练阶段，那么LoRA与预训练权重需要拆开
# if train(False)说明是验证阶段，那么LoRA与预训练权重需要合并
# self.merged是LoRALayer方法，初始是False
should = self.merged if mode else not self.merged
# Let's assume that:
# ⚬ self.weight.data: (384, 128) or (3 * embedding_size, embedding_size)
# ⚬ self.lora_A.data: (4, 128)
# ⚬ self.lora_B.data: (256, 2)
if self.merge_weights and should:
if self.r > 0 and any(self.enable_lora):
#使用分组的conv1d来实现多组的矩阵乘法
delta_w = F.conv1d(
self.lora_A.data.unsqueeze(0), # (4, 128) -> (1, 4, 128)
self.lora_B.data.unsqueeze(-1), # (256, 2) -> (256, 2, 1)
groups=sum(self.enable_lora)
).squeeze(0) # (1, 4, 128) @ (256, 2, 1) -> (1, 256, 128) -> (256, 128)
# -1: W = W - delta_W (unmerge), +1: W = W + delta_W (merge)
#训练阶段需要拆开所以要减
#验证阶段需要合并所以要加
sign = -1 if mode else 1
#self.zero_pad就是扩充到QKV三者都有的shape，这里可以去细看下源码
self.weight.data += sign * self.zero_pad(T(delta_w * self.scaling)) # (256, 128) after zero_pad (384, 128)
self.merged = not mode
以及forward函数
def forward(self, x: torch.Tensor) -> torch.Tensor:
def T(w):
return w.T if self.fan_in_fan_out else w
#forward 函数有了前面的铺垫就比较好理解了，分为两种情况
#推理的时候合并，合并时只需要调用F.linear即可
#训练的时候不合并，需要调用F.linear计算预训练模型链路，也得计算LoRA链路
if self.merged:
return F.linear(x, T(self.weight), bias=self.bias)
else:
# `F.linear` automatically transposes the second argument (T(self.weight) in our case)
result = F.linear(x, T(self.weight), bias=self.bias)
if self.r > 0:
after_A = F.linear(self.lora_dropout(x), self.lora_A)
# For F.conv1d:
# ⚬ input: input tensor of shape (mini-batch, in_channels, iW)
# ⚬ weight: filters of shape (out_channels, in_channels/groups, kW)
# ⚬ groups: split input into groups, in_channels should be divisible by the number of groups. Default: 1
# presumably iW - sequence width/length, kW - kernel width
after_B = F.conv1d(
after_A.transpose(-2, -1),
self.lora_B.unsqueeze(-1),
groups=sum(self.enable_lora)
).transpose(-2, -1)
result += self.zero_pad(after_B) * self.scaling
return result
有了这个MergedLinear的类，我们就可以通过它来替换Attention中生成QKV的三个Linear具体来说如下所示
#原始模型中Attention Block中的Linear
class CausalSelfAttention(nn.Module):
def __init__(self, config: LLaMAConfig) -> None:
super().__init__()
assert config.n_embd % config.n_head == 0
# key, query, value projections for all heads, but in a batch
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
# output projection
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
#带Lora的Attention Block中的Linear
class CausalSelfAttention(llama.CausalSelfAttention):
def __init__(self, config: llama.LLaMAConfig) -> None:
# key, query, value projections for all heads, but in a batch
self.c_attn = MergedLinear(
in_features=config.n_embd,
out_features=3 * config.n_embd,
r=self.lora_config.r,
lora_alpha=self.lora_config.alpha,
lora_dropout=self.lora_config.dropout,
enable_lora=[True, False, True],
fan_in_fan_out = False,
merge_weights=True,
bias=False)
# output projection
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
2.2 对Llama 进行LoRA 微调
有了上述LoRA核心层的代码实现，我们再来看看如何使用LoRA对Llama 进微调，这部分代码在finetune/lora.py
主控逻辑代码如下：
def main(***):
config = LLaMAConfig.from_name("7B")
config.block_size = max_seq_length
checkpoint = torch.load(pretrained_path)
# lora方法就是来实现上述所说的将原始Linear替换成带LoRA的MergedLinear
# lora 是一个 contextmanager方法需要与with语句配合使用
with fabric.init_module(), lora(r=lora_r, alpha=lora_alpha, dropout=lora_dropout, enabled=True):
model = LLaMA(config)
# strict=False because missing keys due to LoRA weights not contained in checkpoint state
model.load_state_dict(checkpoint, strict=False)
#这个方法就是对模型中带lora层的使其requires_grad = True，其他的层使其requires_grad = False
mark_only_lora_as_trainable(model)
#优化器定义
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
model, optimizer = fabric.setup(model, optimizer)
train(fabric, model, optimizer, train_data, val_data, tokenizer_path, out_dir)
# Save the final LoRA checkpoint at the end of training
checkpoint = lora_state_dict(model)
fabric.save(os.path.join(out_dir, "lit-llama-lora-finetuned.pth"), checkpoint)
具体可以看看上述说的lora 方法和mark_only_lora_as_trainable 方法 ，如下所示
@contextmanager
def lora(r, alpha, dropout, enabled: bool = True):
if not enabled:
yield
return
CausalSelfAttention.lora_config = LoRAConfig(r=r, alpha=alpha, dropout=dropout)
causal_self_attention = llama.CausalSelfAttention #保存之前的Linear
llama.CausalSelfAttention = CausalSelfAttention #替换
yield
# when exiting context manager - restore link to original causal self-attention class
llama.CausalSelfAttention = causal_self_attention #处理完with中的语句之后在替换回来
CausalSelfAttention.lora_config = None
def mark_only_lora_as_trainable(model: nn.Module, bias: str = 'none') -> None:
"""Freeze all modules except LoRA's and depending on 'bias' value unfreezes bias weights.
Args:
model: model with LoRA layers #带Lora的层
bias:
``"none"``: all bias weights will be frozen,#所有的bias都不训
``"lora_only"``: only bias weight for LoRA layers will be unfrozen,#只有lora层的bias可训
``"all"``: all bias weights will be unfrozen. #所有bias都能训
Raises:
NotImplementedError: if `bias` not in ["none", "lora_only", "all"]
"""
# freeze all layers except LoRA's
for n, p in model.named_parameters():
if 'lora_' not in n:
p.requires_grad = False
# depending on the `bias` value unfreeze bias weights
if bias == 'none':
return
elif bias == 'all':
for n, p in model.named_parameters():
if 'bias' in n:
p.requires_grad = True
elif bias == 'lora_only':
for m in model.modules():
if isinstance(m, LoRALayer) and \
hasattr(m, 'bias') and \
m.bias is not None:
m.bias.requires_grad = True
else:
raise NotImplementedError
OK 完结撒花～
参考
[1] 苏剑林：梯度视角下的LoRA：简介、分析、猜测及推广
进技术交流群请添加AINLP小助手微信（id: ainlp2)
请备注具体方向+所用到的相关技术点
关于AINLP
AINLP 是一个有趣有AI的自然语言处理社区，专注于 AI、NLP、机器学习、深度学习、推荐算法等相关技术的分享，主题包括LLM、预训练模型、自动生成、文本摘要、智能问答、聊天机器人、机器翻译、知识图谱、推荐系统、计算广告、招聘信息、求职经验分享等，欢迎关注！加技术交流群请添加AINLP小助手微信(id：ainlp2)，备注工作/研究方向+加群目的。
更多AI工具，参考Github-AiBard123，国内AiBard123
