LoRA 工作原理

178 阅读5分钟

问题:随着语言模型容量的不断增大,如从GPT-2,RoBERTa-L到GPT-3,可训练参数量已增长到175B。如果要微调以适应下游任务,如此大的参数量对于时间和显存要求都带来很大的挑战。

为了使模型参数更具效率,有两种常见的策略:

  1. adding apapter layers :参考 smashinggradient.com/2023/04/11/… **** adapter就是固定原有的参数,并添加一些额外参数用于微调。上图中会在原始的transformer block中添加2个adapter,一个在多头注意力后面,另一个这是FFN后面。 loading

    结果:增加的层消耗时间的,并且串联关系,引起Inference Latency

  2. optimizing some forms ´ of the input layer activations, (Directly Optimizing the prompt):(自适应更新和下游任务占比) reduces the sequence length

模型思路:

  • 灵感来源于前人的一些关于 intrinsic dimension 发现,模型是过参数化的,有更小的内在维度
  • 增加模块去学习这个改变量<A,B 模块>

loading

W 冻结不动(GPT-3,stable diffusion)

增加一个可训练矩阵 W\triangle W , W=BA\triangle W = B \cdot A

其中BRd×r,ARr×k,r<<min(d,k)B \in \mathbb{R}^{d \times r}, A \in \mathbb{R}^{r \times k}, r << min(d,k)

h=Wx+Wx=Wx+BAxh = Wx + \triangle W x = W x + BAx

A 使用随机高斯初始化, B 初始化为0

  • 推理的过程中,只需要把改变量放回原模型,就不会有额外的延时(W=W0+BAW = W_0 + BA
  • 切换不同的下游任务的时候,W通过减去 BA得到原始的W0W_0,增加新的 B‘A’去finetune

低秩矩阵分解

在LoRA技术中,秩表示用于分解大矩阵的两个低秩矩阵的维度。

具体来说,假设我们有一个权重矩阵 W ,通过低秩分解,我们将其表示为两个矩阵 A 和 B 的乘积,即:ΔW=A×B

其中, A 的维度是(m,r),B 的维度是(r,n),这里 r 就是秩(Rank)。

选择较小的 r 可以显著减少参数量,从而降低计算和存储成本。

其中, A 和 B 的秩要比 W0 小得多,这样可以显著减少需要调整的参数数量。具体步骤如下:

  1. 预训练模型权重初始化:使用预训练模型的权重矩阵 W0 初始化。
  2. 低秩矩阵初始化:初始化低秩矩阵 A 和 B 。
  3. 微调过程:在微调过程中,只调整低秩矩阵 A 和 B 的参数,而不改变预训练模型的原始权重矩阵 W0 。

优势

  • 降低计算资源需求:通过调整低秩矩阵来适应模型,只需微调较少的参数,大大降低了计算成本和内存需求。
  • 提高适应效率:低秩矩阵分解可以在不显著影响模型性能的情况下,提高微调的效率和速度。
  • 适应不同任务:LoRA可以轻松适应不同的下游任务,只需调整少量的参数即可实现高效的迁移学习。

代码解释

  1. 读取原始模型
import torch
import torch.nn as nn
from transformers import GPT2Model, GPT2Tokenizer

model_name_or_path = "GPT2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name_or_path)
model = GPT2Model.from_pretrained(model_name_or_path)

text = "今天天气非常好"
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)
print(output)

loading

  1. 查看模型的权重参数
for name, param in model.named_parameters():
    print(f"参数名称: {name}, 形状: {param.shape}")

3. 获取其中一层进行lora微调

class LoRAAdapter(nn.Module):
    def __init__(self, original_weight, rank=4):
        super(LoRAAdapter, self).__init__()
        self.rank = rank
        self.A = nn.Parameter(torch.randn(original_weight.size(0), rank))
        self.B = nn.Parameter(torch.randn(rank, original_weight.size(1)))

    def forward(self, W0):
        return W0 + torch.matmul(self.A, self.B)
original_weight = None
for name, param in model.named_parameters():
    if 'attn.c_attn.weight' in name:
        original_weight = param
        break

if original_weight is None:
    raise ValueError("在模型中找不到注意力权重。")

lora_adapter = LoRAAdapter(original_weight)

optimizer = torch.optim.Adam(lora_adapter.parameters(), lr=1e-4)

tokenizer.pad_token = tokenizer.eos_token
if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})

# 数据加载器 (假设已经有一个数据集)
def get_dataloader():
    # 这里使用一个简单的示例数据集
    texts = ["today  i  want to go to the store", "i am fine thank you", "can you help me to go to the store", "I am a student"]
    encodings = tokenizer(texts, return_tensors='pt',padding=True,truncation=True)
    dataset = torch.utils.data.TensorDataset(encodings['input_ids'])
    return torch.utils.data.DataLoader(dataset, batch_size=2)

dataloader = get_dataloader()

# 定义训练过程
def train(model, lora_adapter, dataloader, optimizer, epochs=3):
    model.train()
    for epoch in range(epochs):
        for batch in dataloader:
            inputs = batch[0]
            outputs = model(input_ids=inputs).last_hidden_state

            # 简单的损失函数 (示例)
            loss = outputs.mean()  # 通常你会有一个更复杂的损失函数

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            # 更新模型权重
            with torch.no_grad():
                updated_weight = lora_adapter(original_weight)
                for layer in model.h:
                    layer.attn.c_attn.weight.copy_(updated_weight)

        print(f"第 {epoch + 1}/{epochs} 轮,损失: {loss.item()}")

# 执行微调
train(model, lora_adapter, dataloader, optimizer)

loading

github.com/DWHNicholas…

官方代码

官方在CONV,Embedding,Linear层都实现了lora,这里只贴出Lora在Linear层的实现。全部代码参阅:github.com/microsoft/L…


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)
            if self.r > 0:
                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)

从实现代码中也可以看出,LoRA冻结了PLM的参数,实际需要训练的参数只有lora_A,lora_B,并且,在训练的时候,PLM权重是需要参与计算的,因此LoRA并非是训练高效的