LoRA微调方法详解

leo0308

已于 2023-08-21 19:54:28 修改

阅读量2.9w

点赞数 82

分类专栏：大模型文章标签：人工智能

于 2023-08-15 23:19:48 首次发布

本文链接：https://blog.csdn.net/leo0308/article/details/132308210

版权

大模型专栏收录该内容

24 篇文章

订阅专栏

本文要介绍的是大模型的微调训练方法之一----LoRA。

0 背景

现在大模型非常火爆，大家都在想方设法应用大模型。当前很多大模型虽说可以zero-shot直接使用，但是在具体应用上一般还是微调一下效果更好，也就是常说的finetune。在小模型时代， finetune不是个问题。但大模型时代， finetune是个大问题。这是因为现在的大模型参数动辄10B起，训练的代价非常高昂，即使是finetune也对计算资源有很高要求（finetune只是训练的步数少，对显存等计算资源的占用并没有少)。没个上百G的显存是玩不动的，这对普通人的门槛实在太高了。

那么高效的finetune方式就非常必要了。LoRA就是高效finefune方法的一种。

1 LoRA原理

LoRA论文： LoRA: Low-Rank Adaptation of Large Language Models

在这里插入图片描述
LoRA的原理非常简单，先上一张图，其实从图上已经能清楚地看到大致的原理的。通俗地讲，它的原理是这样的：大模型都是过参数化的，当用于特定任务时，其实只有一小部分参数起主要作用。也就是参数矩阵维度很高，但可以用低维矩阵分解近似。其实这个思想与矩阵特征向量，主成分分析，压缩感知等有异曲同工之妙。

具体做法是，在网络中增加一个旁路结构，旁路是A和B两个矩阵相乘。 A矩阵的维度是dxr, B 矩阵的维度是rxd, 其中r<<d, 一般r取1，2，4，8就够了。那么这个旁路的参数量将远远小于原来网络的参数W。LoRA训练时，我们冻结原来网络的参数W, 只训练旁路参数A和B。由于A和B的参数量远远小于W, 那么训练时需要的显存开销就大约等于推理时的开销。对采用Adam优化器来说，需要的显存就大约相当于全参数finetune的1/3, 极大地减小了训练的代价。

论文中作者的实验也证明了这一点。在GPT-3 175B的finetune中，采用LoRA微调显存的消耗从1.2TB 降低到了350GB, 大约是三分之一

其实采用这种旁路相加的方式，与ResNet的跳连方式也有异曲同工之妙。原网络的参数不变，在旁路上做些微小改变，适应特定新任务。这样就可以让网络基本保持原来的能力，在特定任务上更精进了一步。

值得注意的是， LoRA微调并没有改变原有的预训练参数，只是针对特定任务微调出了新的少量参数，新的这些参数要与原有的预训练参数配合使用（实际使用时，都是把旁路的参数和原来的参数直接合并，也就是参数相加，这样就完全不会增加推理时间）。这是非常方便的，针对不同的任务，都可以训练出自己的LoRA参数，然后与原本的预训练参数结合，做成插件式的应用。这就是最近大火的SD + LoRA。全参数微调一般没这个条件，但LoRA微调还是可以的。目前Civitai上有上万LoRA的模型，并且还在迅速增加。

2 代码详解

LoRA代码： https://github.com/microsoft/LoRA

LoRA原理很简单，代码实现也不复杂。简单地说，在模型实现上，要在特定的模块上加一个旁路，这个旁路就是两个矩阵相乘的形式。这些特定的模块理论上可以是任何模块，目前作者实现的是在Linear, Embeding, Conv, Attention(只改其中的q和v)这些模块上加。

具体实现见：https://github.com/microsoft/LoRA/blob/main/loralib/layers.py

拿其中的Linear做个简单分析吧，其他都是类似的。

class LoRALayer():
    def __init__(
        self, 
        r: int, 
        lora_alpha: int, 
        lora_dropout: float,
        merge_weights: bool,
    ):
        self.r = 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
        # Mark the weight as unmerged
        self.merged = False
        self.merge_weights = merge_weights


class Linear(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.,
        fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
        merge_weights: bool = True,
        **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.fan_in_fan_out = fan_in_fan_out
        # Actual trainable parameters
        if r > 0:
            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))
            self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)))
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
        self.reset_parameters()
        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
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

    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)
        if mode:
            if self.merge_weights and self.merged:
                # Make sure that the weights are not merged
                if self.r > 0:
                    self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
                self.merged = False
        else:
            if self.merge_weights and not self.merged:
                # Merge the weights and mark it
                if self.r > 0:
                    self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
                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
        if self.r > 0 and not self.merged:
            result = F.linear(x, T(self.weight), bias=self.bias)            
            result += (self.lora_dropout(x) @ self.lora_A.transpose(0, 1) @ self.lora_B.transpose(0, 1)) * self.scaling
            return result
        else:
            return F.linear(x, T(self.weight), bias=self.bias)

在Linear层的实现上，多继承了一个LoRALayer, LoRALayer中就是设置了一些参数，最主要的就是上面的讲道的矩阵的秩r了，其他就是一些辅助参数，如控制训练和推理时主路参数和旁路参数是否合并等等。在Linear层中，多定义了A和B两个可训练的参数矩阵，然后在forward中把主路和旁路输出相加，基本上就是完全按照原理来的。

3 使用

实际使用LoRA微调时，也不用自己向上面那样实现了。上面的loralib库已经实现好了，直接使用就好了。具体而言，就是把网络中原来使用nn.Linear用loralib库中的Linear替换就可以了，其他的模块同理。

实际上，还有更简洁的方式，huggingface pert库很贴心地把各种finetune方式都做了集成，更加简单和方便。
huggingface pert库代码： https://github.com/huggingface/peft
官方博客：https://huggingface.co/blog/zh/peft

pert源码解析：https://zhuanlan.zhihu.com/p/646611666