LoRA微调方法详解

本文要介绍的是大模型的微调训练方法之一----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