当红炸子鸡 LoRA,是当代微调 LLMs 的正确姿势?

科技   科技   2024-12-30 22:00   广东  
↑ 点击蓝字 关注极市平台
作者丨CW不要無聊的風格
编辑丨极市平台

极市导读

 

对炼丹界的当红炸子鸡LoRA的「大拷问」!结合源码解析深入了解LoRA。>>加入极市CV技术交流群,走在计算机视觉的最前沿

前言

自 ChatGPT 掀起了大模型(LLM)风潮后,一大波 LLMs(GPT-4, LLaMa, BLOOM, Alpaca, Vicuna, MPT ..) 百花齐放。知识问答、文章撰写、代码编写和纠错、报告策划等等,它们都会,也能够交互式地和你玩文字游戏,甚至还有些很有才的朋友将 LLM 作为交互的接口,同时连接到其它各种模态(e.g. 视觉 & 语音)的模型,从而创造了炸裂的多模态效果,炫~!

这么炫,难免人人都想打造一个自己专属的 LLM(怎么有种回到了小时候玩宠物驯养游戏的赶脚..)。但是,大多数像 CW 这种“平民”玩家,并没有能够玩得起 LLM 的资源(主要是 GPU),别说成百上千亿参数量的模型了,就算是几十亿级别的模型,玩得起的朋友可能也不多。

大多数人对于 LLM 的“亲密度”,可能最多就是拉个开源的 demo 跑下推理过程,得到个“意料之中”的结果,然后很讽刺地自 high 一把:WOW~ 好腻害哟!我们离 AGI 又更近了一步!至于你让他训一个?他会说:呵呵.. 别想太多,洗洗睡吧!

一项技术通常都不会在它诞生之初就得以被广泛使用,和人一样,它也需要机遇。正是现在这种背景,加剧了我们大多数平民在大模型时代炼丹的矛盾。于是,本文的主角 LoRA(Low-Rank Adaptation of Large Language Models) ,一个于2021年就出生的家伙,顺势成为炼丹界的当红炸子鸡,成功出圈。

本文会先介绍 LoRA 的概念与优势、讲述其 motivation 和 以往方法存在的问题,然后以提问的形式从七个方面切入去认识与理解 LoRA(结合源码解析),接着进一步深入思考 LoRA 的一些方面,最后给出一个应用 LoRA 进行微调的例子。对 LoRA 已经有基本了解的帅哥靓女们,可以直接跳到“LoRA 七问”与“进击的 LoRA”这两章。

快给我讲讲 LoRA 是什么

如今快节奏生活下的人们都比较浮躁,你们看我吹水那么多水还没讲 LoRA 到底是什么,肯定已经饥渴难耐。哦?你说你不会,那很好,CW 为你点赞!不过,我也不磨叽,该进入正题了。

LoRA 的全称是 "Low-Rank Adaption", 看到 "low-rank",线性代数玩家们应该会神经反射性地联系到低秩矩阵。Binggo! 这里就是这个意思。你问我 LoRA 的中文名?Em.. 就叫它“低秩(自)适应”吧,虽然英文里没有 "self", 但根据 LoRA 的思想和做法及其带来的效果,它就是自适应的意思。

概括地来说,LoRA 是一项主要用于微调 LLMs 的技术,它额外引入了可训练的低秩分解矩阵,同时固定住预训练权重。这个玩法的重点在于:预训练权重不需训练,因此没有梯度,仅训练低秩矩阵那部分的参数。

有一点 CW 一定要告诉你们:引入的低秩矩阵那部分的参数量比起预训练权重来说,少炒鸡多! 这就意味着,比起全量 fine-tune 的玩法,可训练的参数量少了很多,于是就不需要那么多显存(GPU)资源了。这对于我等平(贫)民来说,简直不要太香了~!

利用 LoRA,我们可以享受到诸多福利,比如以下几点:

  1. 在面对不同的下游任务时,仅需训练参数量很少的低秩矩阵,而预训练权重可以在这些任务之间共享
  2. 省去了预训练权重的梯度和相关的 optimizer states,大大增加了训练效率降低了硬件要求
  3. 训练好的低秩矩阵可以合并(merge)到预训练权重中,多分支结构变为单分支,从而达到没有推理延时的效果;
  4. 与之前的一些参数高效的微调方法(如 Adapter, Prefix-Tuning 等)互不影响,并且可以相互结合

注:参数高效的微调方法 即 PEFT(Parameter-Efficient Fine-Tuning),这类方法仅需微调少量参数(可以是额外引入的),而无需微调预训练模型的所有参数,从而能够降低计算和存储资源。

灵感来源

对于一项技术,CW 往往会好奇它是基于怎样的想法被发明出来的。也就是,发明者的灵感来源究竟源自于哪里。可惜无法亲自采访作者,不然我肯定让他“口若悬河”hhh!没办法咯,我只能通过 paper 来为自己找答案。

CW 发现,作者在 paper 中提到:以往的一些工作表明,模型通常是“过参数化”(over-parametrized)的,它们在优化过程中参数更新的部分通常“驻扎”(reside)在低维子空间中。基于此,作者就顺理成章地提出假设:预训练模型在下游任务中微调而更新参数时,也符合这样的规律。

另外,以往的 PEFT 方法存在一系列问题,如:加大了推理延时、增加了模型深度、限制了输入句长 等,更重要的是,它们大多数都打不过全量 fine-tune,也就是最终训完的模型性能没有全量 fine-tune 来得好。

结合自己的假设与时代背景,作者就搞出了 LoRA,在这种玩法下训出的模型,最终在性能上能和全量 fine-tune 对飙,甚至在一些任务上还更加出色。

以往的 PEFT 跪在哪里

在上一章,CW 浅浅地提到了以往的 PEFT 方法存在的一些问题,如今在本章,我再稍稍展开来谈一下。

在 LoRA 出生之前,比较有代表性的 PEFT 方法主要有额外引入适配下游任务的 adapter layers 和 优化靠近模型输入层的激活(activations) 两种。对于后者,比较有代表性的有 prefix tuning 等。其实,降低下要求,这些方法也算是 good 的,毕竟在一定程度上算是 work 了,只是不 good enough ..

对于引入 adapter layers 这类方法,它们的不足在于:

  1. 新增的 adapter layers 必须串行处理,从而增加了推理延时
  2. 在分布式训练时需要更多的 GPU 同步操作(如 All-Reduce & Broadcast)

至于另一类方法,以 prefix tuning 为例,它们则跪在了:

  1. 该方法本身很难优化
  2. 这种方法需要在模型的输入序列中预留一部分用作可微调的 prompt,从而限制了原始输入文本的句长

LoRA 七问

这一章,CW 会向大家详细说明 LoRA 的玩法,主要从七个方面切入,分别对应以下每一小节的标题,这其实也是我自己在刚接触 LoRA 时所产生的疑问。可以把它们当作一个个 target 去攻破,待全部攻破之后,对 LoRA 应该就算是有一定的理解了。

前四节主要是理论分析,结合了 paper 中的公式和实验结果。后三节的内容则会结合源码解析,这样会有更深刻的认识。

为何可以引入低秩矩阵

作者说他之前看到了一篇 paper: Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning,这篇 paper 的结论是:预训练语言模型在下游任务微调后,权重矩阵其实具有很低的 "intrinsic rank"F-。

【关于 intrinsic rank 的理解】
"intrinsic" 直译是“内在的”、“固有的”,因此我看到有人直接喊 instrinsic rank 为“内在秩”、“固有秩”。(⊙o⊙)… 对于这种叫法,我是觉得很别扭,而且有点不明所以。
CW 觉得,在这里,"intrinsic" 应该理解为“本质上的”、“最具代表性的”会比较恰当。于是,"intrinsic rank" 就应当理解为最能体现数据本质的维度(特征)数目,我们也因此可以美其名曰:“本征秩”。其实,在信号处理里也有相应的概念——intrinsic dimension,它代表能够表示信号的最少特征数,它们所对应的特征是最能体现信号本质的特征。

也就是说,模型在适配下游任务后,权重矩阵的本征秩变得很低,这就代表着其实并不需要这么高的维度数就能够进行表征了,在高维度的权重矩阵中存在冗余。

基于此,作者就自信地认为:模型在微调过程中,权重更新的那部分(参数矩阵)肯定也低秩(low rank)的。

Inspired by this, we hypothesize the updates to the weights also have a low “intrinsic rank” during adaptation.

你问:“权重更新的那部分”具体指什么?

CW 答:模型在微调过程中,权重的变化可以表示为 。其中 就是更新前的权重(在一开始就是预训练权重),而 就是更新的那部分,也就是经过反向传播得到梯度后计算出来的需要更新的量

LoRA 改变了什么

假设 , 由于梯度与权重参数是一对一的, 因此 。如今, 既然认为 的本征秩很低,那么不妨对其做低秩分解:

, 其中 , 并且

这个 就是所谓的 low rank,因为它远小于

由此可知, 经过低秩分解后, 这部分的参数量是远小于预训练权重 的。

按理来说, 是在反向传播阶段才会出现的, 但我们可以将其 “提前拿出来” :让它和 一起做好朋友参与前向过程。这样一来, 反向传播时梯度就只会传导到 这部分, 因其本身就是待更新量, 而 初始是预训练权重, 被固定了, 无需接收梯度。

经过 LoRA 的“洗礼”, 现在如果你喂给模型一个输入 , 前向过程就会表示为:

另外,这里还有两点需要提一下:

  • 低秩矩阵 B,AB,A 的初始化

经过 式的低秩分解后, 为了在一开始让模型的输出保持为原来预训练模型的输出(也就是没有 那部分), 于是将 初始化为全0, 而 则采用随机高斯初始化。

  • 对低秩部分的输出 进行 scale

作者在 paper 中还提到, 对于 这部分, 会乘上一个 scale 系数 。其中, 相对于 保持一个常数倍的关系。作者认为调节这个 大致相当于调节学习率, 于是干脆固定为常数(这样就可以偷懒了 )。

降低了哪部分的显存需求

由于 的参数量远小于 , 因此, 相比于全量 fine-tune 的玩法, LoRA 降低了 optimizer states 这部分对于显存资源的需求。

这是因为 optimizer 对于需要更新的模型参数会保存一份副本, 在全量 fine-tune 的玩法下, 要全量更新, 于是 optimizer 保存的副本参数量为 ; 而我们的小可爱 LoRA 仅需更新 这部分, 所以 optimizer 保存的副本参数量仅为 , 其中 是远小于 的。

另外, 我们可能很容易理所当然地认为 LoRA 对于梯度部分的显存需求也远小于全量 fine-tune, 实际真的是这样吗? 嘿嘿不妨一起来分析下 的梯度是如何计算的。

假设模型经过如公式 所示的前向过程后得到了输出 , 并且我们进一步计算出了损失 ,现在我们是来求 的梯度。根据链式求导法则,易得:

注意 这部分, 它和全量 fine-tune 时是一样的, 这部分梯度的 shape 和权重矩阵的 shape 一致,都是

OMG!这就是说, 实际在计算 的梯度过程中, 所需的显存并没有比全量 fine-tune 来得少, 样也需要算出 shape 为 的梯度矩阵。更尴尬的是, 由于 的存在, 因此在梯度的计算过程中, 所需的显存和计算量甚至比全量 fine-tune 还来得多.. 幸运的是, 在计算完成后, 这个中间状态量所占的显存会被释放掉, 仅需存储 这部分 shape 为 的梯度矩阵。

所以说,对于梯度部分,LoRA 在严格意义上并不能降低其对于显存资源的需求,甚至比起全量 fine-tune 来说计算量还更大了,只不过降低了最终存储的需求。

对模型的哪些部分做低秩分解

可是,在如今的 202x 年代,模型通常有 N 个权重矩阵,那么应该对其中的哪些做低秩分解呢?还是说,应该暴力地、一个不拉地通杀?

对于这个问题,作者选择了偷懒,他仅将 LoRA 应用于 self-attention 层中的 projection matrices(如),而其中的 MLP 模块以及 self-attention 层以外的结构均“不受待见”。

In the Transformer architecture, there are four weight matrices in the self-attention module (Wq, Wk, Wv, Wo) and two in the MLP module.
We limit our study to only adapting the attention weights for downstream tasks and freeze the MLP modules.
We leave the empirical investigation of adapting the MLP layers, LayerNorm layers, and biases to a future work.

作者可能也猜到了你们可能会打破砂锅问到底:应该对 self-attention 层中的哪个或哪几个 projection matrices 应用 LoRA 呢?于是,对于这个问题,他倒是下了番功夫去做实验进行探究。

在实验中,作者以 175B 的 GPT-3 为研究对象,并设置了参数量为 18M 的 budget,也就是应用了 LoRA 部分的可微调参数量不能超过 18M。在这种设置下,当每层仅对 的其中一个应用 LoRA 时,rank 则等于8;而如果每层都对 的其中两个应用 LoRA,则 rank 等于4。

通过上表可以看出,模型更倾向于我们对更多类型的 projection matrices 应用 LoRA(如上表显示,对4个 projection matrices 都应用 LoRA 时效果是最好的),尽管 rank 很低(如上表中最右一列 ), 也足够让 捕获足够的信息。

在代码中如何实现

假设 是一个线性层(Linear Layer),我们一起来看看对其应用 LoRA 是如何实现的。

(麻烦认真看下代码中的注释,谢谢~)

class MergedLinear(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.0,
        enable_lora: List[bool] = [False],
        fan_in_fan_out: bool = False,
        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)

        # enable_lora 是一个布尔类型的列表,用于指示对权重矩阵的哪些“子部分”做低秩分解。
        # 比如 W 是 shape 为 (out_features, in_features) 的矩阵,
        # 那么 enable_lora = [True, False, True] 就表示将 W 在 out_features 这个维度上按序均分成三部分 W1, W2, W3,
        # shape 均为 (out_features // 3, in_features),然后仅对 W1 和 W3 做低秩分解。
        # 其中 W1 的第一个维度取值范围是 [0, out_features // 3),W3 则是 [2 * out_features // 3, out_features)。
        # 同理,若 enable_lora = [True],就表示对整个 W 都做低秩分解。
        if out_features % len(enable_lora) != 0:
            raise ValueError("The length of enable_lora must divide out_features")
        
        self.enable_lora = enable_lora
        self.fan_in_fan_out = fan_in_fan_out

        # Actual trainable parameters
        if r > 0 and any(enable_lora):
            # 仅 enable_lora = True 的部分应用低秩分解,每部分的 low-rank 是 r
            self.lora_A = nn.Linear(in_features, r * sum(enable_lora), bias=False)
            # 注意下这里 B 是用一维的分组卷积实现的
            self.lora_B = nn.Conv1d(
                r * sum(enable_lora),
                out_features // len(enable_lora) * sum(enable_lora),
                kernel_size=1,
                groups=2,
                bias=False,
            )

            # scale 系数,对低秩矩阵的输出(BAx)做缩放
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            # 固定住预训练权重
            self.weight.requires_grad = False

            # Compute the indices
            # 记录权重矩阵中,做了低秩分解的是哪些“子矩阵”
            self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view(len(enable_lora), -1)
            self.lora_ind[enable_lora, :] = True
            self.lora_ind = self.lora_ind.view(-1)

        self.reset_parameters()
        if fan_in_fan_out:
            # fan_in_fan_out 是针对 GPT-2 的 Conv1D 模块的,
            # 该模块和 Linear 的区别就是维度互为转置
            self.weight.data = self.weight.data.T

    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.weight, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B.weight)

以上这个类叫作 MergedLinear, 顾名思义就是低秩分解的部分 可以合并到原来的预训练权重 中。

以上代码中, 比较绕的是与 enable_lora 这个参数相关的内容, 该参数可以用来灵活地指定预训练权重 中,哪些部分要做低秩分解。

关于这个参数的设计由来, 猜想这是因为在某些模型的实现中, Attention 层中的 projection matrix 是用一个共享的线性层实现的(比如 GPT-2, BLOOM, etc.), 而有了这个 enable_lora 参数, 就可以灵活地指定要对这三者中的哪一个做低秩分解。

所有需要进行低秩分解的层都会继承 LoraLayer 这个父类,这个类没有什么特别,也就是设置一些 LoRA 该有的属性:

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.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

        # 是否要禁用低秩分解的部分,如果是,则仅使用预训练权重部分
        self.disable_adapters = False

现在来介绍下 MergedLinear 这个层的前向过程:

  1. 先用预训练权重 对输入 实施前向过程, 得到 ;
  2. 再将输入 喂给低秩分解矩阵 , 得到输出 ;
  3. 接着对 这部分作 "零填充"使其与 的 shape 一致, 并且进行缩放(scale);
  4. 最后再将这部分的结果加回至 中, 如前面的公式
def forward(self, x: torch.Tensor):
    result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
    if self.r > 0:
        after_A = self.lora_A(self.lora_dropout(x))
        after_B = self.lora_B(after_A.transpose(-2-1)).transpose(-2-1)
        result += self.zero_pad(after_B) * self.scaling

        return result

第3点中的“零填充"对应以上代码中的 zero_pad(), 前面 CW 在介绍 enable_lora 这个参数时说过, 由于不一定会对整个预训练权重矩阵做低秩分解, 于是 的 shape 不一定等同于 , 因此需要对前者进行 padding, 使其与后者的 shape 一致, 从而才可让两者进行 element-wise add 。

现在放出这个填充的逻辑:

    def zero_pad(self, x):
        """ 将低秩矩阵的输出 BAx 与原权重矩阵输出 Wx 在维度上对应起来, 维度不足的部分用0填充 """
        result = x.new_zeros((*x.shape[:-1], self.out_features))
        result = result.view(-1, self.out_features)
        # 将 BAx “塞到”与 Wx 相对应的正确位置
        result[:, self.lora_ind] = x.reshape(-1, self.out_features // len(self.enable_lora) * sum(self.enable_lora))

        return result.view((*x.shape[:-1], self.out_features))

怎么做到无推理延时

CW 在前面提到过,LoRA 的优势之一就是推理无延时(相比预训练模型),这是因为低秩分解的部分可以合并至原预训练权重中。比如,这时候模型需要推理,那么你会先调用 model.eval(),这时等同于调用了 model.train(mode=False),接着再将低秩分解部分合并至预训练权重中,过程如下:

    def train(self, mode: bool = True):
        nn.Linear.train(self, mode)
        self.lora_A.train(mode)
        self.lora_B.train(mode)

        # 注:当调用 model.eval() 时就会调用 train(mode=False)
        # 将低秩矩阵 A, B 合并至原权重矩阵 W
        if not mode and self.merge_weights and not self.merged:
            # Merge the weights and mark it
            if self.r > 0 and any(self.enable_lora):
                # \delta_W = BA
                delta_w = (
                    # 这里使用1维卷积将低秩矩阵 A, B 进行“融合”:
                    # A(r * k) 作为输入,r 看作是其 channel,k 看作是空间维度上的大小;
                    # B(d * r * 1) 作为卷积权重,d 是 output channel, r 是 input channel, 1 是 kernel size(注意B本身就是用1维分组卷积实现的)。
                    # 由于是卷积,因此二维的 A 需要增加一维给 mini-batch:r * k -> 1 * r * k。
                    # 卷积后,输入(1 * r * k) -> 输出(1 * d * k)
                    F.conv1d(
                        self.lora_A.weight.data.unsqueeze(0),
                        self.lora_B.weight.data,
                        groups=sum(self.enable_lora),
                    )
                    .squeeze(0)  # 1 * d * k -> d * k
                    .transpose(-2-1)  # d * k -> k * d
                )
                # zero_pad() 是对低秩分解矩阵 \delta_W 进行0填充,因为原权重矩阵 W 中可能有些部分没有进行低秩分解,
                # 从而得到一个和原权重矩阵 W 的 shape 对齐的结果,以便进行加和。k * d -> k * D(假设 D 是原权重矩阵 W 的 out features)
                # 对于原权重矩阵 W 是 Linear 层的情况,fan_in_fan_out = False,于是这里会进行 transpose: k * D -> D * k;
                # 而对于原权重矩阵 W 是 GPT-2 的 Conv1D 的情况,fan_in_fan_out=True,于是不需要 transpose,它的 out features 就是放在第二维的
                # W = W + # \delta_W
                self.weight.data += transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out)
        elif xxx:
            ...

merge 完之后,在进行前向过程时就无需再像上一节展示的那样分步进行,而是一步到位(见以下第二个分支):

def forward(self, x: torch.Tensor):
    # 此部分先省略,下一节再介绍
    if xxx:
        ...
    # 低秩分解部分已合并
    elif self.merged:
        return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
    # 低秩分解部分未合并
    else:
        result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
        if self.r > 0:
            after_A = self.lora_A(self.lora_dropout(x))
            after_B = self.lora_B(after_A.transpose(-2-1)).transpose(-2-1)
            result += self.zero_pad(after_B) * self.scaling

        return result

如何在下游任务中灵活地切换

LoRA 还有一点很吸引人,就是模型在某个下游任务 A 微调后,可以将低秩矩阵那部分的参数解耦出来,还原出预训练权重,从而继续在另一个下游任务 B 中进行微调。

    def train(self, mode: bool = True):
        nn.Linear.train(self, mode)
        self.lora_A.train(mode)
        self.lora_B.train(mode)

        if xxx:
            ...
        # 前一个分支是代表 mode=False,进入该分支说明 mode=True,即调用了 model.train(),
        # 那么当低秩矩阵 A, B 已经合并至原权重矩阵 W 中时,就需要将它们分解出来,以便进行训练(预训练权重 W 无需训练)。
        elif self.merge_weights and self.merged:
            # Make sure that the weights are not merged
            if self.r > 0 and any(self.enable_lora):
                # \delta_W = BA
                delta_w = (
                    F.conv1d(
                        self.lora_A.weight.data.unsqueeze(0),
                        self.lora_B.weight.data,
                        groups=sum(self.enable_lora),
                    )
                    .squeeze(0)
                    .transpose(-2-1)
                )
                # W = W - \delta_W
                self.weight.data -= transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out)
            self.merged = False

还原了预训练权重后,如果你不想使用低秩矩阵那部分的参数,也可以(见以下第一个分支):

    def forward(self, x: torch.Tensor):
        # 当指定不需要使用 adapters 部分(在这里即低秩分解矩阵 \delta_W=BA 这部分),
        # 则将已经合并到预训练权重 W 中的 \delta_W 解耦出来,仅用预训练权重 W 对输入 x 进行前向操作
        if self.disable_adapters:
            if self.r > 0 and self.merged and any(self.enable_lora):
                delta_w = (
                    F.conv1d(
                        self.lora_A.weight.data.unsqueeze(0),
                        self.lora_B.weight.data,
                        groups=sum(self.enable_lora),
                    )
                    .squeeze(0)
                    .transpose(-2-1)
                )
                # W = W - \delta_W
                self.weight.data -= transpose(self.zero_pad(delta_w * self.scaling), not self.fan_in_fan_out)
                self.merged = False
            return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
        # 当使用 adapters 并且低秩分解矩阵 \delta_W=BA 已经合并至预训练权重 W 中,则直接进行前向过程即可。
        elif self.merged:
            ...
        # 当使用 adapters 但低秩分解矩阵 \delta_W=BA 未合并到预训练权重 W 中,则“分步”进行前向过程:
        # 先用预训练权重 W 对输入 x 实施前向过程,得到 Wx;
        # 再将输入 x 喂给低秩分解矩阵 \delta_W=BA,得到 adapters 的输出 (\delta_W)x;
        # 接着对 adapters 部分的输出作零填充使得其与 Wx 的 shape 一致,并且进行缩放(scale);
        # 最后再将这部分的结果加回至 Wx 中。
        else:
            ...

进击的 LoRA

攻破了 LoRA 七问之后,是时候来些更深层次的思想活动了。

Rank r 的设置

一个很直接的问题就是:在实践中,rank 应该设为多少比较合适呢?

作者做了几组实验进行比较,结果发现 rank 可以很低,不超过8就很 OK 了,甚至是1也挺好..

Low Rank 的有效性

看到前面这个实验现象,作者“忍不住”认为: 拥有很低的本征秩(intrinsic rank),增加 rr 并不能使其覆盖到更多有意义的子空间,low rank 万岁!

不过, 他并非是个嘴炮, 对于这一直觉, 他还是仔仔细细地做了个实验去验证的。具体做法就是: 以同样的预训练模型, 分别用 两种 rank 设置去应用 LoRA 进行微调, 然后将训好的低秩矩阵 拿出来做奇异值分解, 得到它们的右奇异单位矩阵 , 最后去比较它们的 top 奇异值向量所张成(span)的子空间的重合程度(使用 Grassmann Distance 度量), 公式表示为(对应 paper 中的公式4):

其中 对应于 的 top-i 奇异值向量的列, 同理。 的取值范围是 , 越大代表两个子空间重合度越高。

根据上图的实验结果显示,两者在最 top 的那些奇异值向量所张成的空间重合度最高,特别是在 top-1 处,这也对前面一节“ 效果也不错”提供了一种解释。

由于在两种 rank 的设置下使用的是同一个预训练模型,并且经过同样的下游训练后,两者在 top 奇异值向量的方向上比较一致(其余方向上则相关度较小),因此这说明 top 奇异值向量所指的方向对于下游任务i 来说是最有用的,而其它方向可能更多地是一些随机噪声的方向,这可能是训练过程中被潜在地积累下来的。

于是乎,low rank 对于 才是正解。

低秩矩阵与预训练权重矩阵的关系

刚接触 LoRA 时, CW 就很好奇训出来的这个 与原来的预训练权重 到底有没有“血缘关系", 它与 的相关度是怎么样的呢?

作者也对这个问题进行了探究, 他将 映射到 维子空间, 得到 ,其中 分别是 左、右奇异向量矩阵。然后计算 的 Frobenius 范数 (后续简称为 范数, 即所有元素的平方和再开方) 。作为对照组, 作者还将 分别映射到了自身和一个随机矩阵的 top-r 奇异值向量空间。

由实验结果可以看出, 预训练权重矩阵 与低秩矩阵 还算是“近亲":比起随机矩阵, 映射到 的子空间后的数值会大一些。

低秩矩阵的优化方向

以上实验结果还揭示出两点:

  • 低秩矩阵在经过下游训练后放大了那些在预训练时没有被强调的向量方向
  • 对于上一点, 较低时放大程度更强 (

咋一看实验结果,可能无法 get 到以上第一点,那么该如何理解呢?

还是得看回上面那张实验结果图, 无论是 还是 的情况, 映射到 维子空间后其 F2 范数都非常小(0.32 & 1.90), 说明这些子空间的向量方向并非是 中重要程度比较高的方向, 而是在预训练时显得不那么重要的、没有被强调的方向。但可以看到, 却不那么小(6.91 & 3.57), 说明在经过下游微调后, 那些原本存在感不强的方向被重视起来了。

由此可知, 在下游训练过程中, 低秩矩阵 并非简单地重复预训练权重矩阵的 top 奇异值向量方向,而是去放大原本在预训练中没有被强调的方向。

至于第二点, 我们分别在 两种情况下计算 , 其中 那一列的值, 这个计算结果衡量了低秩矩阵在第2点中的放大效应。通过计算, 我们可以发现在秩比较的低的情况下 放大效应更强, 那么这又意味着什么呢?

我们有理由认为, 低秩矩阵 包含着大部分与下游任务相关的向量方向(毕竟是朝着下游最优的方向而进行优化的), 于是以上计算结果就意味着适配下游任务的矩阵的本征秩是低秩的。

巧了! 不小心再次证明了 low rank 对于 才是正解~

低秩是万能的吗

CW 在以上多次喊道“low rank 对于 是正解”其实有点过于夸张了。首先,作者的实验场景十分有限,没有在更广泛的 case 上进行验证;其次,我们也不应该无脑地认为以一个很小的 值就能够在所有任务和数据集上 work。

想象下,当下游任务与预训练任务的差异(gap)巨大时(比如在英文上预训练、而在中文上微调),使用很小的 值应该不会有好效果,这时候去微调模型所有的参数(可以令 )应该会得到更好的效果,毕竟中英文的向量空间重合度应该不那么高,我们需要让更多的参数“调头”,转向到适配中文的空间中去。


Example: 在单卡 12G 左右的显存下微调 BLOOM-7B

最后这部分给一个利用 LoRA 进行微调的例子,这个 demo 基于 Huggingface 的 PEFT 库,使用了 LoRA + 8bit 训练,其中 8bit 训练需要安装 bitsandbytes。在单卡的条件下,12G 左右的显存即可玩起 7B 的 BLOOM。

一些杂碎

先来做一些琐碎小事:导入模块、设置数据集相关参数、训练参数 以及 随机种子等。

import gc
import os
import sys
import psutil
import argparse
import threading

import torch
import torch.nn as nn

import numpy as np

from tqdm import tqdm

from torch.utils.data import DataLoader

from datasets import load_dataset
from accelerate import Accelerator
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    default_data_collator,
    get_linear_schedule_with_warmup,
    set_seed,
)

from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_int8_training


def set_seed(seed: int):
    """
    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch` and/or `tf` (if installed).

    Args:
        seed (`int`): The seed to set.
    """


    random.seed(seed)
    np.random.seed(seed)
    
    if is_torch_available():
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        # ^^ safe to call this function even if cuda is not available
    
    if is_tf_available():
        tf.random.set_seed(seed)


def main():
    accelerator = Accelerator()

    dataset_name = "twitter_complaints"
    text_column = "Tweet text"
    label_column = "text_label"
    
    # 样本的最大句长
    max_length = 64

    lr = 1e-3
    batch_size = 8
    num_epochs = 20

    # 设置随机种子(42 yyds!)
    seed = 42
    set_seed(seed)

数据加载和预处理

这里使用的数据集是 RAFT(The Real-world Annotated Few-shot Tasks) 任务中的 Twitter Complaints,有50个训练样本和3399个测试样本。

下面给出数据加载和预处理的逻辑,代码本身简单明了,无需啰嗦。

''' Datset and Dataloader '''

    dataset = load_dataset("ought/raft", dataset_name, cache_dir=args.data_cache_dir)
    classes = [k.replace("_"" "for k in dataset["train"].features["Label"].names]
    dataset = dataset.map(
        lambda x: {"text_label": [classes[label] for label in x["Label"]]},
        batched=True,
        num_proc=4
    )

    # Preprocessing

    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, cache_dir=args.model_cache_dir)

    def preprocess_function(examples):
        # 注:这里的 batch size 并非训练时的 batch size,
        # 在这个数据预处理的过程中,也是批处理的,所以这里也有个 batch size 的概念
        batch_size = len(examples[text_column])

        # Add prompt 'Label' to input text
        inputs = [f"{text_column} : {x} Label : " for x in examples[text_column]]
        targets = [str(x) for x in examples[label_column]]

        model_inputs = tokenizer(inputs)
        labels = tokenizer(targets)

        # 依次处理每个样本
        for i in range(batch_size):
            sample_input_ids = model_inputs["input_ids"][i]
            label_input_ids = labels["input_ids"][i] + [tokenizer.pad_token_id]

            # 将输入文本(model_inputs)与标签(labels)“对齐”(设置成一样),然后将标签中对应输入文本的部分设为 -100,
            # 这样在计算 loss 时就不会计算这部分,仅计算真实标签文本的那部分。
            # Add label text to input text
            model_inputs["input_ids"][i] = sample_input_ids + label_input_ids
            # Let the label value which correspond to the input text word to be -100
            labels["input_ids"][i] = [-100] * len(sample_input_ids) + label_input_ids
            model_inputs["attention_mask"][i] = [1] * len(model_inputs["input_ids"][i])

        # Put pad tokens at the front of the inputs, and truncate to 'max_length'
        for i in range(batch_size):
            sample_input_ids = model_inputs["input_ids"][i]
            label_input_ids = labels["input_ids"][i]

            pad_length = max_length - len(sample_input_ids)
            labels["input_ids"][i] = [-100] * pad_length + label_input_ids
            model_inputs["input_ids"][i] = [tokenizer.pad_token_id] * pad_length + sample_input_ids
            model_inputs["attention_mask"][i] = [0] * pad_length + model_inputs["attention_mask"][i]

            # To tensor
            model_inputs["input_ids"][i] = torch.tensor(model_inputs["input_ids"][i][:max_length])
            model_inputs["attention_mask"][i] = torch.tensor(model_inputs["attention_mask"][i][:max_length])
            labels["input_ids"][i] = torch.tensor(labels["input_ids"][i][:max_length])

        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    def test_preprocess_function(examples):
        batch_size = len(examples[text_column])

        inputs = [f"{text_column} : {x} Label : " for x in examples[text_column]]
        model_inputs = tokenizer(inputs)

        for i in range(batch_size):
            sample_input_ids = model_inputs["input_ids"][i]

            pad_length = max_length - len(sample_input_ids)
            model_inputs["input_ids"][i] = [tokenizer.pad_token_id] * pad_length + sample_input_ids
            model_inputs["attention_mask"][i] = [0] * pad_length + model_inputs["attention_mask"][i]

            # To tensor
            model_inputs["input_ids"][i] = torch.tensor(model_inputs["input_ids"][i][:max_length])
            model_inputs["attention_mask"][i] = torch.tensor(model_inputs["attention_mask"][i][:max_length])

        return model_inputs

    with accelerator.main_process_first():
        processed_datasets = dataset.map(
            preprocess_function,
            batched=True,
            num_proc=4,
            remove_columns=dataset["train"].column_names,
            load_from_cache_file=True,
            desc="Running tokenizer on dataset",
        )
    accelerator.wait_for_everyone()
    train_dataset = processed_datasets["train"]

    with accelerator.main_process_first():
        processed_datasets = dataset.map(
            test_preprocess_function,
            batched=True,
            num_proc=4,
            remove_columns=dataset["train"].column_names,
            load_from_cache_file=False,
            desc="Running tokenizer on dataset",
        )
    eval_dataset = processed_datasets["train"]
    test_dataset = processed_datasets["test"]

    # Dataloaders
    train_dataloader = DataLoader(
        train_dataset, shuffle=True, collate_fn=default_data_collator,
        batch_size=batch_size, pin_memory=True, num_workers=4
    )
    eval_dataloader = DataLoader(
        eval_dataset, collate_fn=default_data_collator,
        batch_size=batch_size, pin_memory=True, num_workers=4
    )
    test_dataloader = DataLoader(
        test_dataset, collate_fn=default_data_collator,
        batch_size=batch_size, pin_memory=True, num_workers=4
    )

    print(f"The 1st train batch sample: {next(iter(train_dataloader))}\n")

Model, Optimizer & Lr scheduler

必备套装:模型、优化器、学习率的调度。

''' Model, Optimizer, Lr Scheduler '''

    # creating model
    model = AutoModelForCausalLM.from_pretrained(
        args.model_name_or_path,
        cache_dir=args.model_cache_dir,
        load_in_8bit=args.load_in_8bit,
        device_map='auto'  # A device map needs to be passed to run convert models into mixed-int8 format
    )

    ''' Post-processing on the model, includes:
        1- Cast the layernorm in fp32;
        2- making output embedding layer require grads;
        3- Anable gradient checkpointing for memory efficiency;
        4- Add the upcasting of the lm head to fp32
    '''

    model = prepare_model_for_int8_training(model)

    # 配置 LoRA 的一些参数
    peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1)
    # 对模型应用 LoRA
    model = get_peft_model(model, peft_config)

    # 打印出可训练的参数个数
    model.print_trainable_parameters()

    # optimizer
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)

    # lr scheduler
    lr_scheduler = get_linear_schedule_with_warmup(
        optimizer=optimizer,
        num_warmup_steps=0,
        num_training_steps=(len(train_dataloader) * num_epochs),
    )

    model, train_dataloader, eval_dataloader, test_dataloader, optimizer, lr_scheduler = accelerator.prepare(
        model, train_dataloader, eval_dataloader, test_dataloader, optimizer, lr_scheduler
    )
    accelerator.print(f"Model: {model}\n")

Preparation for 8bit Training

这一节进一步解析上一节中的 model = prepare_model_for_int8_training(model) 这部分,它是为了让训练过程更稳定,以便得到更好的效果,下面来看看其具体做了些什么。

def prepare_model_for_int8_training(
    model, output_embedding_layer_name="lm_head", use_gradient_checkpointing=True, layer_norm_names=["layer_norm"]
)
:

    r"""
    This method wrapps the entire protocol for preparing a model before running a training. This includes:
        1- Cast the layernorm in fp32 2- making output embedding layer require grads 3- Add the upcasting of the lm
        head to fp32

    Args:
        model, (`transformers.PreTrainedModel`):
            The loaded model from `transformers`
    """

    loaded_in_8bit = getattr(model, "is_loaded_in_8bit"False)

    # 1. 固定预训练的权重;
    # 2. 将 Layer Norm 的参数转换为 fp32,这是为了训练的稳定性  
    for name, param in model.named_parameters():
        # freeze base model's layers
        param.requires_grad = False

        if loaded_in_8bit:
            # cast layer norm in fp32 for stability for 8bit models
            if param.ndim == 1 and any(layer_norm_name in name for layer_norm_name in layer_norm_names):
                param.data = param.data.to(torch.float32)

    # 让 Embedding 层接受梯度,通过对 Embedding 层注册 forward hook 实现,
    # forward hook 的内容会在模型前向过程完成后被调用。
    # 通过以下可以看到,这里 hook 的内容是使 Embedding 层的输出接受梯度,
    # 从而梯度可以传导到 Embedding 层。
    if loaded_in_8bit and use_gradient_checkpointing:
        # For backward compatibility
        if hasattr(model, "enable_input_require_grads"):
            model.enable_input_require_grads()
        else:

            def make_inputs_require_grad(module, input, output):
                output.requires_grad_(True)

            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)

        # enable gradient checkpointing for memory efficiency
        # 对前向过程的中间激活部分的梯度做优化,以优化内存所需。
        model.gradient_checkpointing_enable()

    # 将模型头部的输出转换为 fp32 以稳定训练
    if hasattr(model, output_embedding_layer_name):
        output_embedding_layer = getattr(model, output_embedding_layer_name)
        input_dtype = output_embedding_layer.weight.dtype

        class CastOutputToFloat(torch.nn.Sequential):
            r"""
            Manually cast to the expected dtype of the lm_head as sometimes there is a final layer norm that is casted
            in fp32

            """


            def forward(self, x):
                # 这里之所以要先将输入(x)转换成该层参数的精度(dtype)是因为上一层可能是 Layer Norm,
                # 而由上可知,我们对 Layer Norm 的输出精度转换成了 fp32,因此在这种情况下,就需要先将
                # 上一层的输出(也就是该层的输入 x)先转成与该层参数同样的精度。
                return super().forward(x.to(input_dtype)).to(torch.float32)

        setattr(model, output_embedding_layer_name, CastOutputToFloat(output_embedding_layer))

    return model

通过以上的源码实现并结合 CW 的注释,可以知道,这部分主要做了以下四件事情:

  1. LayerNorm 的参数转换成 fp32 类型;
  2. Embedding 层的输出接收梯度,从而使得梯度得以传递至该层;
  3. 优化前向过程产生的中间激活部分的梯度,以减少显存消耗;
  4. 将模型头部(Head)的输出精度转换成 fp32 类型

PEFT Model

在这一节,CW 引领大家来看看从普通的 model 转换成 peft model:_model = get_peft_model(model, peft_config)_ 是怎么做的,以下针对 BLOOM 模型的情况进行解析。

def get_peft_model(model, peft_config):
    """
    Returns a Peft model object from a model and a config.

    Args:
        model ([`transformers.PreTrainedModel`]): Model to be wrapped.
        peft_config ([`PeftConfig`]): Configuration object containing the parameters of the Peft model.
    """


    model_config = model.config.to_dict()
    peft_config.base_model_name_or_path = model.__dict__.get("name_or_path"None)
    
    if peft_config.task_type not in MODEL_TYPE_TO_PEFT_MODEL_MAPPING.keys():
        peft_config = _prepare_lora_config(peft_config, model_config)
        return PeftModel(model, peft_config)
    
    if not isinstance(peft_config, PromptLearningConfig):
        # BLOOM 会进入到这个分支
        peft_config = _prepare_lora_config(peft_config, model_config)
    else:
        peft_config = _prepare_prompt_learning_config(peft_config, model_config)

    # 在我们这个例子里,peft_config.task_type 是 CAUSAL_LM,
    # MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type] 则是 PeftModelForCausalLM,
    # 它是 PeftModel 的子类,它就是在原模型基础上对目标模块做了 LoRA 转换的结果
    return MODEL_TYPE_TO_PEFT_MODEL_MAPPING[peft_config.task_type](model, peft_config)

进一步来看看 peft_config = _prepare_lora_config(peft_config, model_config) 这里面的实现,它决定了要对模型的哪些模块应用 LoRA 。

def _prepare_lora_config(peft_config, model_config):
    if peft_config.target_modules is None:
        if model_config["model_type"not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING:
            raise ValueError("Please specify `target_modules` in `peft_config`")
        # 设置需要进行 LoRA 转换的目标模块,通常是 Attention 层中的一个或几个映射矩阵(Linear Layer)
        # 对于 BLOOM,这里返回的是 ["query_key_value"],对应的是其模型实现中 BloomAttention 中的 QKV 映射矩阵
        peft_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config["model_type"]]

    if len(peft_config.target_modules) == 1:
        # 这个仅对 GPT-2 里用到的 Conv1D 有效
        peft_config.fan_in_fan_out = True
        # 这三个值分别代表 Q,K,V 映射矩阵是否要应用 LoRA
        # 对于 BLOOM 来说,这里仅对 Q, V 的映射矩阵做转换,而 K 不做。
        peft_config.enable_lora = [TrueFalseTrue]
    
    if peft_config.inference_mode:
        # 如果是推理模式,则将低秩矩阵 A, B 合并到 Linear 层原来的权重 W 中
        peft_config.merge_weights = True
        
    return peft_config

结合 CW 在上面代码中的注释可知,对于 BLOOM,peft_config.target_modules 是 ["query_key_value"],这对应的是其子模块 BloomAttention 中的 Q, K, V 映射矩阵:

class BloomAttention(nn.Module):
    def __init__(self, config: BloomConfig):
        super().__init__()
        
        # 省略部分
        ...

        self.hidden_size = config.hidden_size
        self.num_heads = config.n_head
        self.head_dim = self.hidden_size // self.num_heads
        self.split_size = self.hidden_size
        self.hidden_dropout = config.hidden_dropout

        # 省略部分
        ...

        # ["query_key_value"] 指的就是这个模块
        self.query_key_value = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=True)
        self.dense = nn.Linear(self.hidden_size, self.hidden_size)
        self.attention_dropout = nn.Dropout(config.attention_dropout)

这个 target_modules 也支持自定义,只要和模型实现里的关键字匹配得上就行。

训练

其实就是常规的训练迭代,只不过这里的特色是利用了 TorchTracemalloc 上下文管理器,它可以方便地计算出 GPU 和 CPU 的消耗(以 MB 计)。

    for epoch in range(num_epochs):
        with TorchTracemalloc() as tracemalloc:
            model.train()

            total_loss = 0
            for step, batch in enumerate(tqdm(train_dataloader)):
                # Forward
                outputs = model(**batch)
                loss = outputs.loss
                total_loss += loss.detach().float()

                # Backward
                accelerator.backward(loss)
                optimizer.step()
                lr_scheduler.step()

                optimizer.zero_grad()

                if step % 3 == 0:
                    accelerator.print(f"epoch {epoch + 1}\t step {step + 1}\t loss {loss.item()}")
            
            epoch_loss = total_loss / len(train_dataloader)
            epoch_ppl = torch.exp(epoch_loss)
            accelerator.print(f"[Epoch{epoch + 1}]\t total loss: {epoch_loss}\t perplexity: {epoch_ppl}\n")
                
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
        accelerator.print("GPU Memory before entering the train : {}".format(b2mb(tracemalloc.begin)))
        accelerator.print("GPU Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.used))
        accelerator.print("GPU Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.peaked))
        accelerator.print(
            "GPU Total Peak Memory consumed during the train (max): {}\n".format(
                tracemalloc.peaked + b2mb(tracemalloc.begin)
            )
        )

        accelerator.print("CPU Memory before entering the train : {}".format(b2mb(tracemalloc.cpu_begin)))
        accelerator.print("CPU Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.cpu_used))
        accelerator.print("CPU Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.cpu_peaked))
        accelerator.print(
            "CPU Total Peak Memory consumed during the train (max): {}\n".format(
                tracemalloc.cpu_peaked + b2mb(tracemalloc.cpu_begin)
            )
        )

        train_epoch_loss = total_loss / len(eval_dataloader)
        train_ppl = torch.exp(train_epoch_loss)
        accelerator.print(f"{epoch=}{train_ppl=} {train_epoch_loss=}\n")

顺便秀一波训练期间 GPU 和 CPU 的资源消耗情况(以下单位均为 MB):

某个 epoch 的训练资源消耗

哦?你说你好奇 TorchTracemalloc 是如何实现的?OK,CW 也不吝啬,这就为您献上:

def b2mb(x):
    """ Converting Bytes to Megabytes. """
    return int(x / 2**20)


class TorchTracemalloc:
    """ This context manager is used to track the peak memory usage of the process. """

    def __enter__(self):
        gc.collect()
        torch.cuda.empty_cache()
        # Reset the peak gauge to zero
        torch.cuda.reset_max_memory_allocated()

        # 返回当前的显存占用
        self.begin = torch.cuda.memory_allocated()
        self.process = psutil.Process()

        self.cpu_begin = self.cpu_mem_used()

        self.peak_monitoring = True
        peak_monitor_thread = threading.Thread(target=self.peak_monitor_func)
        peak_monitor_thread.daemon = True
        peak_monitor_thread.start()

        return self

    def cpu_mem_used(self):
        """Get resident set size memory for the current process"""
        return self.process.memory_info().rss

    def peak_monitor_func(self):
        self.cpu_peak = -1

        while True:
            self.cpu_peak = max(self.cpu_mem_used(), self.cpu_peak)

            # can't sleep or will not catch the peak right (this comment is here on purpose)
            # time.sleep(0.001) # 1msec

            if not self.peak_monitoring:
                break

    def __exit__(self, *exc):
        self.peak_monitoring = False

        gc.collect()
        torch.cuda.empty_cache()

        self.end = torch.cuda.memory_allocated()
        self.peak = torch.cuda.max_memory_allocated()

        self.used = b2mb(self.end - self.begin)
        self.peaked = b2mb(self.peak - self.begin)

        self.cpu_end = self.cpu_mem_used()
        self.cpu_used = b2mb(self.cpu_end - self.cpu_begin)
        self.cpu_peaked = b2mb(self.cpu_peak - self.cpu_begin)

评估

评估与训练的玩法基本类似,只不过前向过程需要调用的是模型的 generate() 方法,而非 forward(),前者是 auto-regressive 的方式。

        model.eval()

        eval_preds = []
        with TorchTracemalloc() as tracemalloc:
            for batch in tqdm(eval_dataloader):
                batch = {k: v for k, v in batch.items() if k != "labels"}

                with torch.no_grad():
                    # 注:推理过程用的是 auto-regressive 的方式,调用的是模型的 generate() 方法
                    outputs = accelerator.unwrap_model(model).generate(**batch, max_new_tokens=10)
                outputs = accelerator.pad_across_processes(outputs, dim=1, pad_index=tokenizer.pad_token_id)
                
                preds = accelerator.gather(outputs)
                # The part before 'max_length' belongs to prompts
                preds = preds[:, max_length:].detach().cpu().numpy()
                # 'skip_special_tokens=True' will ignore thoses special tokens(e.g. pad token)
                eval_preds.extend(tokenizer.batch_decode(preds, skip_special_tokens=True))

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
        accelerator.print("GPU Memory before entering the eval : {}".format(b2mb(tracemalloc.begin)))
        accelerator.print("GPU Memory consumed at the end of the eval (end-begin): {}".format(tracemalloc.used))
        accelerator.print("GPU Peak Memory consumed during the eval (max-begin): {}".format(tracemalloc.peaked))
        accelerator.print(
            "GPU Total Peak Memory consumed during the eval (max): {}\n".format(
                tracemalloc.peaked + b2mb(tracemalloc.begin)
            )
        )

        accelerator.print("CPU Memory before entering the eval : {}".format(b2mb(tracemalloc.cpu_begin)))
        accelerator.print("CPU Memory consumed at the end of the eval (end-begin): {}".format(tracemalloc.cpu_used))
        accelerator.print("CPU Peak Memory consumed during the eval (max-begin): {}".format(tracemalloc.cpu_peaked))
        accelerator.print(
            "CPU Total Peak Memory consumed during the eval (max): {}\n".format(
                tracemalloc.cpu_peaked + b2mb(tracemalloc.cpu_begin)
            )
        )

        assert len(eval_preds) == len(dataset["train"][label_column]), \
            f"{len(eval_preds)} != {len(dataset['train'][label_column])}"

        correct = total = 0        
        for pred, true in zip(eval_preds, dataset["train"][label_column]):
            if pred.strip() == true.strip():
                correct += 1
            total += 1

        accuracy = correct / total * 100
        accelerator.print(f"{accuracy=}\n")
        accelerator.print(f"Pred of the first 10 samples:\n {eval_preds[:10]=}\n")
        accelerator.print(f"Truth of the first 10 samples:\n {dataset['train'][label_column][:10]=}\n")

推理期间,GPU 和 CPU 的消耗情况如下(以下单位均为 MB):

某个 epoch 训练后,推理的资源消耗


End

LoRA 作为当今大模型时代最火的技术之一,是否算得上是微调 LLMs 的正确姿势由你们决定。比起正确与否,合不合适才是最重要的。于我而言,只是觉得它好玩而不是无聊的风格而已~

公众号后台回复“数据集”获取100+深度学习各方向资源整理

极市干货

技术专栏:多模态大模型超详细解读专栏搞懂Tranformer系列大视觉模型 (LVM) 解读扩散模型系列极市直播
技术综述:小目标检测那点事大模型面试八股含答案万字长文!人体姿态估计(HPE)入门教程

点击阅读原文进入CV社区

收获更多技术干货

极市平台
为计算机视觉开发者提供全流程算法开发训练平台,以及大咖技术分享、社区交流、竞赛实践等丰富的内容与服务。
 最新文章