(本次博文的LoRA代码主要基于lit-llama工程)
1. lit-llama工程
- lit-llama是一个大语言模型的工程,可用于语言模型的训练、测试等
- 优点:很简洁,代码也很容易上手,基于lightning框架(lightning和pytorch的关系,有点像keras和tensorflow)
- 项目地址: https://github.com/Lightning-AI/lit-llama
2. LoRA部分
这里我们围绕3个问题,对LoRA进行一个大体上的理解:什么是LoRA?LoRA有什么用?LoRA是怎么实现的?
-
什么是LoRA?
LoRA可以认为是一种方法,或者是一种结构或模块
-
LoRA有什么用?
- 通过引入LoRA模块,可以快速地对大语言模型进行微调(传统的微调是非常费时费力的)
-
LoRA怎么实现的?
- 通过冻结原来预训练模型的权重,并引入LoRA模块(可以简单分解为两个全连接层),训练的时候只更新LoRA模块的权重即可
- 通过冻结原来预训练模型的权重,并引入LoRA模块(可以简单分解为两个全连接层),训练的时候只更新LoRA模块的权重即可
3. 代码部分
-
首先看lit-llama工程中LLAMA的模型定义,在lit-llama/model.py文件,通过不同的配置配置选项设置不同size的LLAMA模型,如下图所示:
-
看到模型定义中有个Block模块(其他就是正常的transformer的decoder部分),Block的定义在model.py同一个文件,如下图所示,其中有一个CausalSelfAttention模块:
-
CausalSelfAttention中定义了self.c_attn,这个模块其实就是输入经过Wq, Wk, Wv三个权重矩阵,输入经过这个self.c_attn之后,就得到了Q、K、V(这里不清楚的同学去复习一下自注意力的计算过程)
-
LoRA在代码中的实现,就是在定义模型的时候,替换了CausalSelfAttention中的self.c_attn,为什么呢?因为LoRA默认是作用于attention中的Q和V权重矩阵,即Wq和Wv,不对Wk做操作。
-
查看工程中lit_llama/LoRA.py文件,也定义了一个CausalSelfAttention类。对比两个类,如我们上面所说,发现LoRA中定义的self.c_attn被替换为了MergedLinear
-
具体的计算过程,可以看LoRA.py文件,其实已经在代码中讲的很清楚了,代码中也举例进行了说明,这里已经假设了如下的参数:
- batch_size: 64
- context_length: 64 (这个可以理解为输入文本或者序列的长度)
- embedding_size: 128
- r: 2
-
我们直接跳到MergedLinear类的前向推理部分,如下图所示:
-
LoRA计算的过程如下图所示:
4. LoRA模型文件
-
打开之前指令微调过的LoRA权重文件,如下图所示:
-
由于之前选用LLAMA-7B的配置,其中r = 8, embedding_size = 4096,可以看到有:
- lora_A的shape为:16 x 4096(1个是Wq的8x4096,1个是Wv的8x4096,所以总的是16x4096)
- lora_B的shape为:8192 x 8 (其中一部分为4096x8,用于跟Q的部分做运算,另一部分也为4096x8,用于跟V的部分做运算,所以总的是8192x8)
- LLAMA-7B有32个head,所以对应也有32个LoRA模块。
结束。