大模型训练——PEFT与LORA介绍

吴笛横吹

大模型训练——PEFT与LORA介绍0.简介1.LORA原理介绍2.补充资料：低显存学习方法3.PEFT对LORA的实现0.简介朋友们好，我是练习NLP两年半的算法工程师常鸿宇，今天介绍一下大规模模型的轻量级训练技术LORA，以及相关模块PEFT。Parameter-EfficientFine-Tuning(PEFT)，是huggingface开发的一个python工具，项目地址：https://github.com/huggingface/peft其可以很方便地实现将普通的HF模型变成用于支持轻量级fine-tune的模型，使用非常便捷，目前支持4种策略，分别是：LoRAORAOW-RANKADAPTATIONOFLARGELANGUAGEMODELSPrefixTuningrefix-Tuning:OptimizingContinuousPromptsforGeneration,P-Tuningv2romptTuningCanBeComparabletoFine-tuningUniversallyAcrossScalesandTasksP-Tuning:GPTUnderstands,TooPromptTuning:ThePowerofScaleforParameter-EfficientPromptTuning今天要介绍的，是其中之一，也是最近比较热门的LORA(LOW-RANKADAPTATIONOFLARGELANGUAGEMODELS)。1.LORA原理介绍LORA的论文写的比较难读懂，但是其原理其实并不复杂。简单理解一下，就是在模型的Linear层，的旁边，增加一个“旁支”，这个“旁支”的作用，就是代替原有的参数矩阵W进行训练。结合上图，我们来直观地理解一下这个过程，输入xxx，具有维度ddd，举个例子，在普通的transformer模型中，这个xxx可能是embedding的输出，也有可能是上一层transformerlayer的输出，而ddd一般就是768或者1024。按照原本的路线，它应该只走左边的部分，也就是原有的模型部分。而在LORA的策略下，增加了右侧的“旁支”，也就是先用一个Linear层A，将数据从ddd维降到rrr，这个rrr也就是LORA的秩，是LORA中最重要的一个超参数。一般会远远小于ddd，尤其是对于现在的大模型，ddd已经不止是768或者1024，例如LLaMA-7B，每一层transformer有32个head，这样一来ddd就达到了4096.接着再用第二个Linear层B，将数据从rrr变回ddd维。最后再将左右两部分的结果相加融合，就得到了输出的hidden_state。对于左右两个部分，右侧看起来像是左侧原有矩阵WWW的分解，将参数量从d∗dd*dd∗d变成了d∗r+d∗rd*r+d*rd∗r+d∗r，在r0.0:self.lora_dropout=nn.Dropout(p=lora_dropout)else:self.lora_dropout=lambdax:x#Marktheweightasunmergedself.merged=Falseself.merge_weights=merge_weightsself.disable_adapters=False12345678910111213141516171819然后就要讲到上文中所提到的Linear类，也就是Lora的具体实现，它同时继承了nn.Linear和LoraLayer。classLinear(nn.Linear,LoraLayer):#Loraimplementedinadenselayerdef__init__(self,in_features:int,out_features:int,r:int=0,lora_alpha:int=1,lora_dropout:float=0.0,fan_in_fan_out:bool=False,#SetthistoTrueifthelayertoreplacestoresweightlike(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#Actualtrainableparametersifr>0:self.lora_A=nn.Linear(in_features,r,bias=False)self.lora_B=nn.Linear(r,out_features,bias=False)self.scaling=self.lora_alpha/self.r#Freezingthepre-trainedweightmatrixself.weight.requires_grad=Falseself.reset_parameters()iffan_in_fan_out:self.weight.data=self.weight.data.T123456789101112131415161718192021222324252627在构造方法中，除了对各个超参数进行配置之外，还对所有参数进行了初始化，定义如下：defreset_parameters(self):nn.Linear.reset_parameters(self)ifhasattr(self,"lora_A"):#initializeAthesamewayasthedefaultfornn.LinearandBtozeronn.init.kaiming_uniform_(self.lora_A.weight,a=math.sqrt(5))nn.init.zeros_(self.lora_B.weight)123456其中lora的A矩阵采用了kaiming初始化，是Xavier初始化针对非线性激活函数的一种优化；B矩阵采用了零初始化，以确保在初始状态ΔW=BA\DeltaW=BAΔW=BA为零。（值得注意的是在LORA的论文中，A采用的是Gaussian初始化）。对于train和eval方法，放在一起介绍，它主要是需要对merge状态进行记录：deftrain(self,mode:bool=True):nn.Linear.train(self,mode)self.lora_A.train(mode)self.lora_B.train(mode)ifnotmodeandself.merge_weightsandnotself.merged:#Mergetheweightsandmarkitifself.r>0:self.weight.data+=(transpose(self.lora_B.weight@self.lora_A.weight,self.fan_in_fan_out)*self.scaling)self.merged=Trueelifself.merge_weightsandself.merged:#Makesurethattheweightsarenotmergedifself.r>0:self.weight.data-=(transpose(self.lora_B.weight@self.lora_A.weight,self.fan_in_fan_out)*self.scaling)self.merged=Falsedefeval(self):nn.Linear.eval(self)self.lora_A.eval()self.lora_B.eval()1234567891011121314151617181920212223首先对于新定义的这个Linear层，其本身继承了torch.nn.Linear，所以需要调用nn.Linear.train(self,mode)来控制一下自身原本参数的状态，并且此外它加入了lora_A和lora_B两部分额外的参数，这两部分本质上也是nn.Linear，也需要控制状态。然后主要来理解一下merge_weights是在做什么，也就是看train中的if分支，notmode说明是eval模式，而self.merge_weights在上文中有介绍，是配置文件中的，意思是评估时是否需要将lora部分的weight加到linear层原本的weight中，notself.merged是状态的记录，也就是说，如果设置了需要融合，而当前状态没有融合的话，就把lora部分的参数scale之后加上去，并且更新self.merged状态；在elif分支中，是为了在训练的过程中，确保linear本身的weights是没有经过融合过的（理论上这一步应该是在eval之后的下一轮train的第一个step触发）。至于为什么是在train中涉及merge_weights，其实在torch的源码中，nn.Linear.eval()实际上是调用了nn.Linear.train(mode=False)，所以这里train方法中的merge_weigths，实际上是在eval中也发挥作用的。forward中也是类似的原理，正常情况下训练过程应该是走elif的分支：defforward(self,x:torch.Tensor):ifself.disable_adapters:ifself.r>0andself.merged:self.weight.data-=(transpose(self.lora_B.weight@self.lora_A.weight,self.fan_in_fan_out)*self.scaling)self.merged=FalsereturnF.linear(x,transpose(self.weight,self.fan_in_fan_out),bias=self.bias)elifself.r>0andnotself.merged:result=F.linear(x,transpose(self.weight,self.fan_in_fan_out),bias=self.bias)ifself.r>0:result+=self.lora_B(self.lora_A(self.lora_dropout(x)))*self.scalingreturnresultelse:returnF.linear(x,transpose(self.weight,self.fan_in_fan_out),bias=self.bias)12345678910111213141516在了解了这些基本原理之后，就可以类似地去实现更多更加灵活的功能了，例如对transformer的某些层增加lora，而其余的层保持不变等。以上就是关于LORA的代码实现介绍，在实际的PEFT模块中，还包含了更多更详细完备的设置，本文只是对基本原理和过程进行了介绍，其中包含了部分个人理解，如果错误，还请指出。如果本文对你的学习和工作有所帮助，记得留下一个免费的赞，我们下期再见。