LoRA[1]: 微调大模型的一种轻量级方法

最新推荐文章于 2024-08-13 21:25:33 发布

莫叶何竹

最新推荐文章于 2024-08-13 21:25:33 发布

阅读量962

点赞数 28

分类专栏： LLM 文章标签： lora LORA PEFT lora实现

本文链接：https://blog.csdn.net/weixin_40779727/article/details/141091190

版权

LLM 专栏收录该内容

3 篇文章 0 订阅

订阅专栏

	info
paper	https://arxiv.org/abs/2106.09685
github	https://github.com/microsoft/LoRA
个人博客位置	http://myhz0606.com/article/lora
创作日期：	2023-3-10, 2024-8-10（补充代码实现）

背景

当需要在特定任务上提高通用大模型的泛化能力时，微调是一个非常常用的方法。然而，在通用领域预训练的大模型模型参数非常庞大，以GPT-3为例，参数数量高达175B，这意味着没有多少公司具备微调如此大规模模型的钞能力。因此，轻量化微调技术在今年引起了学术界和工业界的广泛关注和研究。一些代表性的轻量化微调技术包括Adapter和P-Tuning等。在本文中，我们将介绍一种从低维本征维度的角度出发的技术，即LoRA（Low-Rank Adaption），以解决这个问题。

方法大意

作者的思路来源于2020 Facebook Aghajanyan发表的一篇论文4，论文的主要观点是：常见的预训练模型有非常低的本征维度。通俗的讲就是说存在一种低维重参数化方式，其在微调时与全参数空间一样有效。受此启发LoRA相对原本架构它增加了一个旁路，使输入向较小的子空间进行随机投影。微调过程时仅更新旁路的较为轻量的权重A,B，取代更新原有权重W。

用公式描述：

原来的fintune

前向过程	$h = W x$
反向传播过程	根据 $\frac{\partial L}{\partial h}\frac{\partial h}{\partial W}$ 来更新 $W$ （ $L$ 为loss）

对于LoRA

前向过程	$h = W x + B A x$
反向传播过程	根据 $\frac{\partial L}{\partial h}\frac{\partial h}{\partial A}$ , $\frac{\partial L}{\partial h}\frac{\partial h}{\partial B}$ 来更新权重参数 $A, B$ ,其中 $\in \mathbb{R}^{d \times r}, B \in \mathbb{R}^{r \times d}$ ，其中 $\ll d$

训练完毕可以将 $W$ 和 $B, A$ 进行合并 $W_{out}=W + BA$ ，不会增加前向推理的时间。

通过这个操作微调的参数从 $d^2$ 降低到了 $2 r d$ ，大大提升了微调效率。论文通过这个方法仅用35M的参数量微调了GPT-3,并取得了接近全量微调（参数量350GB）的效果。

在这里插入图片描述

思路扩展

LoRA通过增加一个旁路，使用低维重参数化方式来降低微调参数。虽然这使得微调参数降低近10000倍，但训练效率并没有显著提升，因为根据链式求导规则，对原有支路计算梯度仍然需要相同的计算量。为了解决这个问题，后续的工作LST[5]对其进行了优化。

LoRA的训练思路同样被应用于图像任务，例如在Stable-Diffusion微调中得到了广泛的应用[3]。

Lora简易实现

下面笔者将从dinov2出发来实现用Lora。

目标：fix dinov2-vit-base模型的预训练权重，在dinov2的qkv 层中引入lora 可训练参数。（注意：dinov2的qkv层的module type为nn.Linear）

step1: 找到dinov2 qkv层的模型参数

首先通过torchhub加载dinov2模型

dinov2 = torch.hub.load(
            "facebookresearch/dinov2",
            model="dinov2_vitb14",
            source="github"
        )

通过dinov2.named_module可以找到目标层（qkv）

for name, module in dinov2.named_modules():
    if isinstance(module, nn.Linear) and 'qkv' in name:
        print(">>>", name, module.__class__.__name__)

>>> blocks.0.attn.qkv Linear
>>> blocks.1.attn.qkv Linear
>>> blocks.2.attn.qkv Linear
>>> blocks.3.attn.qkv Linear
>>> blocks.4.attn.qkv Linear
>>> blocks.5.attn.qkv Linear
>>> blocks.6.attn.qkv Linear
>>> blocks.7.attn.qkv Linear
>>> blocks.8.attn.qkv Linear
>>> blocks.9.attn.qkv Linear
>>> blocks.10.attn.qkv Linear
>>> blocks.11.attn.qkv Linear

step2: 将qkv层引入lora可训练参数

找到想要添加的层后，我们要将原始的nn.Linear module改为自定义的Linear module

参考前文阐述的Lora原理，我们先构建自定义的 Linear 类

为了扩展，先定义一个LoraLayer基类

class LoraLayer(nn.Module):
    def __init__(self):
        super(LoraLayer, self).__init__()

再定义LoraConfig，用于设置Lora的参数

from dataclasses import dataclass

@dataclass
class LoraConfig:
    rank: int  # lora的rank
    inject_module_name: Union[str, List[str], None] = None  # Lora control的module名称
    inject_module_type: type = nn.Linear # Lora control的module类别
    dropout_p: float = 0.0  # Lora 层的droupout ratio

自定义的Linear类继承LoraLayer

class Linear(LoraLayer):
    def __init__(
            self,
            base_layer: nn.Module,  # 需要control的module
            lora_config: LoraConfig,
            **kwargs
    ):
        super(Linear, self).__init__()
        self.base_layer = base_layer
        # 继承base_layer的一些基础属性
        self.factory_kwargs = {'device': base_layer.weight.device, 'dtype': base_layer.weight.dtype}
        self.in_features = base_layer.in_features
        self.out_features = base_layer.out_features
        self.lora_config = lora_config
        self.merged = False  # Lora参数书否和原始参数合并
			
        # 定义Lora层的参数
        self.lora_A = nn.Linear(self.in_features, self.lora_config.rank, bias=False, **self.factory_kwargs)
        self.lora_B = nn.Linear(self.lora_config.rank, self.out_features, bias=False, **self.factory_kwargs)
        if self.lora_config.dropout_p > 0.0: # 是否对Lora层引入dropout
            self.lora_dropout_layer = nn.Dropout(p=self.lora_config.dropout_p)
        else:
            self.lora_dropout_layer = nn.Identity()

        # 初始化Lora参数
        self.lora_A.weight.data = torch.randn(self.lora_A.weight.data.size())
        self.lora_B.weight.data = torch.zeros(self.lora_B.weight.data.size())
		     
        # 将Lora层的参数设置为可接受梯度
        self.lora_A.weight.requires_grad = True
        self.lora_B.weight.requires_grad = True

    def forward(self, x: torch.Tensor):
        if not self.merged:
            # 当没有合并权重时的forward
            ori_output = self.base_layer(x)
            delta_output = self.lora_B(self.lora_A(self.lora_dropout_layer(x)))
            output = ori_output + delta_output
        else:
            # 合并权重后的forward和base_layer的一致
            output = self.base_layer(x)
        return output

定义好lora-based Linear 层后，下面的重点是如何将原来的层用替换为新定义的层

step3: 将qkv层替换为Lora-based qkv层

我们知道qkv层是Attention层的一个子模块, 比如 module name=”blocks.0.attn.qkv”的qkv层，我们可以通过get_submodule 方法拿到它的parent

parent = dinov2.get_submodule("blocks.0.attn")
print(parent)

MemEffAttention(
  (qkv): Linear(in_features=768, out_features=2304, bias=True)
  (attn_drop): Dropout(p=0.0, inplace=False)
  (proj): Linear(in_features=768, out_features=768, bias=True)
  (proj_drop): Dropout(p=0.0, inplace=False)
)

随后可以通过setattr 方法将原来的qkv层改为lora-based qkv层

target_layer = Linear(
    base_layer=dinov2.get_submodule("blocks.0.attn.qkv"),
    lora_config=LoraConfig(8, 'qkv', nn.Linear)
)
setattr(parent, "qkv", target_layer)
print(parent)

我们再打印一下parent

MemEffAttention(
  (qkv): Linear(
    (base_layer): Linear(in_features=768, out_features=2304, bias=True)
    (lora_A): Linear(in_features=768, out_features=8, bias=False)
    (lora_B): Linear(in_features=8, out_features=2304, bias=False)
    (lora_dropout_layer): Identity()
  )
  (attn_drop): Dropout(p=0.0, inplace=False)
  (proj): Linear(in_features=768, out_features=768, bias=True)
  (proj_drop): Dropout(p=0.0, inplace=False)
)

可以看到qkv层已经成功替换。

下面对所有的qkv层都进行同样的操作

def statistic_parameter(model):
    trainable_params_count = 0
    freeze_params_count = 0
    for name, param in model.named_parameters():
        if param.requires_grad:
            trainable_params_count += param.numel()
        else:
            freeze_params_count += param.numel()
    return {"trainable_params_count": trainable_params_count, "freeze_params_count": freeze_params_count}

def _get_submodules(model, key):
        parent = model.get_submodule(".".join(key.split(".")[:-1]))
        target_name = key.split(".")[-1]
        target = model.get_submodule(key)
        return parent, target, target_name 
    
    
def _replace_module(parent, child_name, new_module):
    setattr(parent, child_name, new_module) 

    
def _create_module(lora_config: LoraConfig, target: nn.Module):
    base_layer = target
    return Linear(base_layer, lora_config=lora_config)
    

def _create_and_replace(
        lora_config: LoraConfig,
        target: nn.Module,
        target_name: str,
        parent: nn.Module,
    ) -> None:
        new_module = _create_module(lora_config, target=target)
        _replace_module(parent, target_name, new_module)
        

def inject_adapter(
    model: nn.Module, 
    lora_config: LoraConfig, 
):
    inject_module_type = lora_config.inject_module_type
    inject_module_name = lora_config.inject_module_name
    
    for name, module in dinov2.named_modules():
        if isinstance(module, inject_module_type):
            if inject_module_name is None:  # replace all linear module to lora-linear
                parent, target, target_name = _get_submodules(model, name)
                _create_and_replace(lora_config=lora_config, target=target, target_name=target_name, parent=parent)
            else:
                if isinstance(inject_module_name, str):
                    inject_module_name = [inject_module_name]
                assert isinstance(inject_module_name, list)
                if sum([True if i in name else False for i in inject_module_name]) > 0:
                    parent, target, target_name = _get_submodules(model, name)
                    _create_and_replace(lora_config=lora_config, target=target, target_name=target_name, parent=parent)

# initial model
dinov2 = torch.hub.load(
    "facebookresearch/dinov2",
    model="dinov2_vitb14",
    source="github"
)

dinov2.eval()
# ========
# freeze backbone
for p in dinov2.parameters():
    p.requires_grad = False
# ========
print(f">>> before inject lora")
print(statistic_parameter(dinov2))
inject_adapter(dinov2, lora_config=LoraConfig(8, 'qkv', nn.Linear))
print(f">>> after inject lora")
print(statistic_parameter(dinov2))

Output:

>>> before inject lora
{'trainable_params_count': 0, 'freeze_params_count': 86580480}
>>> after inject lora
{'trainable_params_count': 294912, 'freeze_params_count': 86580480}

可以看到以成功添加lora的训练参数。

综上对Lora的代码实现做了简单介绍，如有错误，欢迎指出交流。

代码能力强的同学，可以直接阅读huggingface peft相关的源码。

参考资料

[1] LoRA paper

[2] LoRA github

[3] 应用LoRa微调stable diffusion

[4] Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning

[5] 《LST: Ladder Side-Tuning for Parameter and Memory Efficient Transfer Learning》

[6] https://huggingface.co/docs/diffusers/main/en

莫叶何竹

关注

28
点赞
踩
24

收藏

觉得还不错? 一键收藏
0
评论
LoRA[1]: 微调大模型的一种轻量级方法

当需要在特定任务上提高通用大模型的泛化能力时，微调是一个非常常用的方法。然而，在通用领域预训练的大模型模型参数非常庞大，以GPT-3为例，参数数量高达175B，这意味着没有多少公司具备微调如此大规模模型的钞能力。因此，轻量化微调技术在今年引起了学术界和工业界的广泛关注和研究。一些代表性的轻量化微调技术包括Adapter和P-Tuning等。在本文中，我们将介绍一种从低维本征维度的角度出发的技术，即LoRA（ow-ankdaption），以解决这个问题。
复制链接

扫一扫

专栏目录