【学习】手写一个lora微调算法~

手写一个LoRA算法

参考视频： https://www.bilibili.com/video/BV1ZZ421T7XJ/?share_source=copy_web&vd_source=2ff04f0cc70237133af9bcd6be27a652

---

初始化LoRA Linear类

首先 LoRA 的原理就是在预训练模型的权重矩阵上添加两个 LoRA_A 和 LoRA_B 矩阵组合输出嘛，由于 LoRA_A 和 LoRA_B 矩阵可以通过秩 $r$ 来调整规模，因此训练起来参数量就比 FT 小了很多。

所以实际上在训练的时候，我们是通过冻结原模型参数矩阵，来学 LoRA_A 和 LoRA_B 矩阵的参数来实现的。那么就需要将原模型里的 Linear 层给替换成封装了 LoRA 模块的 LoRA Linear 层。

LoraLinear 类的初始化如下：

class LoraLinear(nn.Module):
    def _init_(self, in_features, out_features, merge, rank=16, alpha=16, dropout=0.5):
        super(LoraLinear, self)._init_()
        self.in_features = in_features
        self.out_features = out_features
        self.merge = merge
        self.rank = rank
        self.alpha = alpha
        self.dropout = dropout

in_features：输入维度，比如词嵌入维度是768，这里就是768；

out_features：输出维度;

merge：控制是否合并 LoRA 和原始权重;

rank：LoRA 的秩，越小压缩程度越大，控制 LoRA 矩阵参数量的主要因素;

alpha：缩放系数，用于调节 LoRA 矩阵参数的影响程度;

dropout：防止过拟合的 dropout 率;

        self.linear = nn.Linear(in_features, out_features)
        if rank > 0:
            # 构建权重参数
            self.lora_b = nn.Parameter(torch.zeros(out_features, rank))
            self.lora_a = nn.Parameter(torch.zeros(rank, in_features))
            self.scale = self.alpha / rank
            self.linear.weight.requires_grad = False

        if dropout > 0:
            self.dropout = nn.Dropout(self.dropout)
        else:
            self.dropout = nn.Identity()
        
        self.initial_weights()
        
    def initial_weights(self):
        nn.init.kaiming_uniform_(self.lora_a, a=math.sqrt(5))

创建 LoRA 的 A 和 B 矩阵：A (rank × in_features) 和 B (out_features × rank)；

计算缩放因子 scale = alpha / rank；

使用 Kaiming 初始化来初始化 LoRA 矩阵 A，这里只需要初始化一个权重矩阵，另一个 B 矩阵初始化为零矩阵。这样训练开始时 A @ B = zeros，随着训练进行会慢慢调整 B 矩阵的参数值；

前向过程

    def forward(self, x):
        if self.rank > 0 and self.merge:
            # 可以看到这边lora_b是out_features * rank的矩阵，lora_a是rank * in_features的矩阵，乘完后是out_features * in_features的矩阵
            # 那么weight就是out_features * in_features的矩阵，刚好和linear.weight的形状一样
            # 而wx+b，因此x就是in_features * batch_size的矩阵
            output = F.linear(x, weight=self.linear.weight + self.lora_b @ self.lora_a * self.scale, bias=self.linear.bias)
            output = self.dropout(output)
            return output
        else:
            return self.linear(x)

这里就是用 LoRA 的 A 和 B 矩阵乘上原始权重矩阵的计算过程来在前向中参与进计算过程。在训练的时候就会反向更新 LoRA 两个矩阵的参数。

最简单的实现过程就是这样，假若要将 LoRA 模块包装到具体的一个模型进行微调，那么可以把 base_model 里的 module 里面的线性层替换成 LoRA Linear 层，然后再冻结预训练权重在下游任务上进行训练。可以直接利用 hugging face 的 Trainer 类开一个 Trainer 直接训。