【转载】使用 LoRA(低秩适配)微调大语言模型的实用技巧

原文地址:Practical Tips for Finetuning LLMs Using LoRA (Low-Rank Adaptation),by Sebastian Raschka, on 2023-12-09

使用 LoRA(低秩适配)微调大语言模型的实用技巧

——从上百次实验中总结的经验

LoRA(低秩适配,论文见https://arxiv.org/abs/2106.09685)是高效训练定制大语言模型时应用最广泛、效果最显著的技术之一。对于关注开源大语言模型的从业者而言,这是一项值得深入掌握的核心技术。

本文的实验基于开源项目 Lit-GPT(https://github.com/Lightning-AI/lit-gpt ,我与Lightning AI 的同事共同维护该项目),相关详细实验记录见 https://lightning.ai/pages/community/lora-insights/ 。本篇《Ahead of AI》专栏文章旨在梳理我从大量实验中提炼出的核心经验,同时解答该领域最常见的若干问题。如果你正从事定制大语言模型的微调工作,希望这些经验能帮你少走弯路(原文“the long run”为双关,既指“长期来看”,也暗合模型训练的运行过程)。

简而言之,本文讨论的核心结论如下:

  • 尽管大语言模型训练(以及基于GPU的模型训练普遍)存在固有随机性,但多次重复实验的结果一致性非常高。
  • QLoRA 是一种在显存受限场景下值得权衡的方案:它能节省33%的显存,但训练耗时会增加39%。
  • 微调大语言模型时,优化器的选择并非核心影响因素。单独使用SGD效果欠佳,但无论是使用AdamW、带动量调度的SGD,还是带动量调度的AdamW,最终效果差异极小。
  • 人们常说Adam是显存占用很高的优化器,因为它为每个模型参数维护两个额外状态量,但这对大语言模型的峰值显存需求影响很小——因为显存主要被大规模矩阵乘法占用,而非存储额外参数。
  • 对于固定数据集,多轮迭代训练(即多epoch训练)未必有益,反而往往会导致效果下降,原因大概率是过拟合。
  • 使用LoRA时,应将其应用到模型的所有层,而非仅作用于键(Key)和值(Value)矩阵,才能最大化模型性能。
  • 调整LoRA的秩(rank)至关重要,同时也要选择合适的alpha值;一个实用的经验法则是将alpha设为秩的2倍。
  • 70亿参数的模型可以在单张14GB显存的GPU上高效完成微调,仅需数小时。
  • 仅靠固定数据集,无法让大语言模型在所有基准测试上都表现优异;要实现这一点需要多样化的数据源,或者说LoRA本身可能并非解决该问题的最佳工具。

此外,我还将解答关于LoRA的10个常见问题:
Q1:数据集的重要性有多高?
Q2:LoRA适用于领域适配吗?
Q3:如何选择最优的秩(rank)?
Q4:必须在所有层都启用LoRA吗?
Q5:如何避免过拟合?
Q6:其他优化器的效果如何?
Q7:还有哪些因素会影响显存占用?
Q8:与全量微调、RLHF相比效果如何?
Q9:多个LoRA权重可以合并吗?
Q10:逐层最优秩适配(Layer-wise Optimal Rank Adaptation)效果如何?

(在上一期《Ahead of AI》中我提到,如果大家感兴趣,我会写一篇更基础的入门文章,从零开始实现LoRA的代码。从反馈来看大家的需求很高,我计划后续再发布一篇相关文章。本篇文章则聚焦于LoRA实践中的宏观思路与核心结论,采用自上而下的视角展开。)

LoRA 简介

大语言模型参数量庞大,受GPU显存限制,训练时更新所有权重的成本极高。

举例来说,假设一个70亿参数的大语言模型,其权重矩阵为W(实际模型的参数分布在许多层的不同矩阵中,这里简化为单个权重矩阵说明)。在反向传播过程中,我们会得到一个ΔW矩阵,它记录了为最小化训练损失,原始权重需要更新的幅度。

权重更新公式如下:
$$W_{\text{updated}} = W + \Delta W$$

如果权重矩阵W包含70亿参数,那么权重更新矩阵ΔW同样包含70亿参数,计算ΔW矩阵会消耗极高的算力与显存。

LoRA方法由论文https://arxiv.org/abs/2106.09685提出,它将权重更新量ΔW分解为低秩表示。准确来说,LoRA不需要显式计算ΔW,而是在训练中直接学习ΔW的分解形式,这正是其节省显存的核心原理,如下图所示:

figure01

如图所示,ΔW的分解是指用两个更小的LoRA矩阵A和B来表示大矩阵ΔW。若A的行数与ΔW相同,B的列数与ΔW相同,则分解可表示为ΔW = AB(AB为矩阵A与B的乘积)。

这种方法能节省多少显存?这取决于超参数秩(r)的大小。举个例子:假设ΔW有10000行、20000列,共存储2亿个参数。如果取r=8,那么A为10000行×8列,B为8行×20000列,总参数量为 10000×8 + 8×20000 = 24万个参数,仅为原来2亿的约1/830。

当然,A和B无法完整捕获ΔW包含的全部信息,但这正是LoRA的设计初衷。我们的假设是:预训练阶段的权重矩阵W需要是满秩的大矩阵,才能学习到预训练数据集中的全部知识;但在微调大语言模型时,我们不需要更新所有权重,仅用远少于ΔW的参数量就能捕获适配任务的核心信息——因此我们通过AB的低秩形式来完成权重更新。

1. LoRA 实验的一致性

通过多次重复LoRA微调实验我发现,尽管大语言模型训练(以及GPU训练普遍)存在固有随机性,但不同轮次的基准测试结果一致性非常高。这为后续的对比实验提供了可靠的基础。

figure02

(注:以上结果采用默认设置,秩r=8。实验细节可参考我的另一篇文章:https://lightning.ai/pages/community/lora-insights/

2. QLoRA 的算力-显存权衡

QLoRA(量化低秩适配,论文见https://arxiv.org/abs/2305.14314)是一种进一步降低微调显存占用的技术。在反向传播过程中,QLoRA将预训练权重量化为4位精度,并使用分页优化器处理显存峰值。

实验证实,使用QLoRA可以节省33%的GPU显存,但代价是训练耗时增加39%——这是由于QLoRA需要额外的预训练权重量化与反量化操作。

采用16位bfloat16精度的常规LoRA:

  • 训练时长:1.85小时
  • 显存占用:21.33 GB

采用4位归一化浮点数的QLoRA:

  • 训练时长:2.79小时
  • 显存占用:14.18 GB

此外我发现,QLoRA对模型效果的影响微乎其微。在常见的GPU显存瓶颈场景下,QLoRA是常规LoRA训练的可行替代方案。

figure03

3. 学习率调度器

学习率调度器会在训练过程中逐步降低学习率,以优化收敛效果,避免越过损失函数的最小值点。

余弦退火(Cosine Annealing)是一种遵循余弦曲线调整学习率的调度策略:训练开始时学习率较高,之后平滑下降,以余弦函数的趋势趋近于零。常用的是半周期变体,即训练全程只完成半个余弦周期,如下图所示:

figure04

在实验中,我为LoRA微调脚本加入了余弦退火调度器,观察到它能显著提升SGD的训练效果;但对Adam和AdamW优化器的影响较小,几乎没有明显差异。

figure05

SGD相比Adam的潜在优势将在下一节讨论。

《Ahead of AI》是读者支持的出版物。若想接收新文章并支持我的创作,欢迎成为免费或付费订阅者。

4. Adam 与 SGD 对比

Adam和AdamW优化器在深度学习中仍被广泛使用,但在大模型场景下它们的显存占用很高。原因在于Adam类优化器会为每个模型参数维护两个滑动平均值:梯度的一阶矩(均值)和二阶矩(未中心化方差)。换句话说,Adam会为每个模型参数在显存中额外存储两个值。对于70亿参数的模型,训练时就要额外跟踪140亿个参数。

SGD优化器在训练中不需要跟踪额外参数,因此一个自然的问题是:将Adam替换为SGD,对大语言模型训练的峰值显存需求有多大改善?

在我的实验中,使用AdamW+默认LoRA设置(r=8)训练70亿参数的Llama 2模型,GPU显存占用为14.18 GB;换成SGD训练同一模型,显存占用为14.15 GB。换言之,仅节省了0.03 GB,效果微乎其微。

为什么显存节省这么少?因为使用LoRA时,可训练参数的总量本身就很小。例如当r=8时,70亿参数的Llama 2模型中,可训练的LoRA参数仅为4,194,304个,远小于总参数量6,738,415,616。

单看数字,400多万可训练参数似乎不少,但计算一下:4,194,304 × 2 × 16 bit = 134.22 Mbit = 16.78 MB。(实际观测到0.03 GB = 30 MB的差异,是因为存储和复制优化器状态还有额外开销。)其中“2”是Adam额外存储的参数数量,16位是模型权重的默认精度。

figure06

但在后续实验中,当我把LoRA的秩r提升到256时,Adam和SGD优化器的显存差异就变得明显了:

  • AdamW:17.86 GB
  • SGD:14.46 GB

结论:当LoRA的秩r较小时,把Adam换成SGD的价值不大;但当r较大时,这种替换是值得考虑的。

5. 多轮训练(多Epoch)

在传统深度学习中,我们通常会对训练集进行多轮迭代——每完整遍历一次训练集称为一个epoch(轮次)。例如训练卷积神经网络时,跑上百个epoch很常见。那么多轮训练对指令微调也同样有效吗?

当我把Alpaca指令微调数据集(https://github.com/tatsu-lab/stanford_alpaca)的迭代次数翻倍(相当于2个训练轮次)时,我观察到模型性能出现了下降。

figure07

结论是:多轮训练对指令微调未必有益,反而可能降低效果。在只有1000个样本的LIMA数据集上我也观察到了同样的现象。性能下降大概率是过拟合加剧导致的,这一点还需要更多研究验证。

6. 在更多层启用 LoRA

前面表格中的实验,仅在部分权重矩阵上启用了LoRA——也就是每个Transformer层中的键(Key)和值(Value)权重矩阵。除此之外,我们还可以在查询(Query)权重矩阵、投影层、多头注意力块之间的其他线性层,以及线性输出层上都启用LoRA。

figure08

对于70亿参数的Llama 2模型,若在所有这些额外层都启用LoRA,可训练参数量会变为原来的5倍:从4,194,304增加到20,277,248。对应的显存需求也会上升(从14.18 GB增至16.62 GB),但模型性能会得到显著提升。

figure09

不过我的实验存在局限:只探索了两种设置——(1)仅在查询和值权重矩阵启用LoRA;(2)在所有层启用LoRA。未来可以探索更多组合,比如仅在投影层启用LoRA是否真的有增益,这是一个值得研究的问题。

7. 平衡 LoRA 超参数:秩(R)与 Alpha

正如LoRA原论文(https://arxiv.org/abs/2106.09685)所述,在前向传播时,LoRA权重会通过一个额外的缩放系数叠加到预训练权重上。缩放系数由秩参数r和另一个超参数α(alpha)共同决定,公式如下:

scaling = alpha / r
weight += (lora_B @ lora_A) * scaling

从上面的公式可以看出,alpha越大,LoRA权重的影响就越强。

之前的实验使用r=8、alpha=16,对应2倍缩放。在大语言模型中使用LoRA时,alpha取r的2倍是通用经验法则,但我好奇当r更大时这个法则是否依然成立。

实验表明,“alpha = 2×秩”确实是效果很好的取值。不过在本次实验的模型+数据集组合下,当r=256、alpha=128(对应0.5倍缩放)时,性能甚至更优。

figure10

(我还测试了r=32、r=64、r=128、r=512的情况,为了图表清晰未列出,其中r=256效果最佳。)

将alpha设为r的2倍通常能得到较好的结果,但尝试不同的比例也不会有坏处。

8. 单GPU训练70亿参数模型

LoRA最核心的价值之一,就是让我们能在单张GPU上完成70亿参数大语言模型的微调。在本次实验中,采用最优设置的QLoRA(r=256,alpha=512)+ AdamW,训练5万个样本(Alpaca数据集)仅需约3小时(A100显卡),显存占用17.86 GB。

figure11

在文章余下部分,我将解答大家可能关心的其他问题。

常见问题解答

Q1:数据集的重要性有多高?

数据集至关重要。我在实验中使用了包含5万个训练样本的Alpaca数据集,选择它是因为该数据集非常流行;而测试不同数据集超出了本文范围,因为文章篇幅已经很长。

但需要注意:Alpaca是通过调用旧版ChatGPT生成的合成数据集,以今天的标准来看可能并非最优选择。

数据质量非常关键。比如今年6月我讨论过的LIMA数据集(https://magazine.sebastianraschka.com/p/ahead-of-ai-9-llm-tuning-and-dataset),是一个经过精心筛选的数据集,仅包含1000个样本。

根据论文https://arxiv.org/abs/2305.11206,在LIMA上微调的650亿参数Llama模型,效果明显优于在Alpaca上微调的同规模模型。

figure12

使用最优配置(r=256,alpha=512)在LIMA上微调,得到的效果与样本量50倍于它的Alpaca数据集相当,甚至更优。

figure13

Q2:LoRA适用于领域适配吗?

很遗憾,我对这个问题没有非常确定的答案。一般来说,模型的知识主要来自预训练数据集;指令微调更多是引导大语言模型学会遵循指令。

但值得注意的是,如果显存受限,LoRA也可以用于在领域特定数据集上对已有的预训练大语言模型做继续预训练。

我的实验中包含两个算术基准测试(详见https://lightning.ai/pages/community/lora-insights/ ),结果显示LoRA微调后的模型在这两项上的表现远差于预训练基座模型。我的假设是:因为Alpaca数据集中没有对应的算术样例,模型“遗忘”了算术能力。至于模型是完全丢失了这部分知识,还是只是无法响应相关指令,还需要进一步研究。但这里可以得到一个结论:微调大语言模型时,最好包含所有你关心的任务类型的样例。

Q3:如何选择最优的秩(rank)?

很遗憾,目前没有很好的经验法则可以直接选出合适的r,它是一个需要针对每个模型、每个数据集单独探索的超参数。

我的推测是:r过大会导致更严重的过拟合;而r太小则可能无法覆盖数据集中多样化的任务。换句话说,数据集中的任务越多样,r的取值就应该越大。比如,如果只需要模型做基础的两位数算术,很小的r可能就足够了。但这只是假设,还需要更多实验验证。

Q4:必须在所有层都启用LoRA吗?

我只探索了两种设置:(1)仅在查询和值权重矩阵启用LoRA;(2)在所有层启用LoRA。未来可以探索更多组合,比如仅在投影层启用LoRA是否真的有增益,这是一个值得研究的方向。

figure14

举例来说,仅考虑lora_querylora_keylora_valuelora_projectionlora_mlplora_head这几个开关,就有2⁶=64种组合可以探索。这是未来研究一个很有意思的方向。

Q5:如何避免过拟合?

一般来说,r越大可训练参数越多,过拟合风险就越高。如果模型出现过拟合,首先可以尝试降低r,或者扩大数据集规模。此外,还可以提高AdamW或SGD优化器的权重衰减率,也可以考虑增大LoRA层的dropout值。

我在实验中没有深入探索LoRA的dropout参数(使用了固定的0.05 dropout率),这是未来值得研究的方向。

Q6:其他优化器的效果如何?

未来还有很多有意思的大语言模型优化器值得探索,其中一个是今年5月发布的Sophia(论文见https://arxiv.org/abs/2305.14342)。

Sophia是一种二阶优化算法,在Adam和AdamW占据主流的大语言模型领域,它的表现很有吸引力。论文显示,相比Adam,Sophia的训练速度快2倍,且训练出的模型效果更好。简而言之,Sophia是通过梯度曲率而非梯度方差来对梯度做归一化,这是它与Adam的核心区别。

Q7:还有哪些因素会影响显存占用?

除了精度和量化设置、模型大小、批次大小、可训练LoRA参数量之外,数据集本身也会影响显存占用。

注意Llama 2的上下文窗口(block size)是4096个token。大语言模型的上下文窗口决定了它一次能处理的最大序列长度;但训练序列越短,由于未来token掩码的存在,显存节省就越明显。

例如Alpaca数据集的序列长度相对较短,最大长度为1304个token。

figure15

当我用最大长度2048个token的其他数据集实验时,显存占用从17.86 GB上升到了26.96 GB。

Q8:与全量微调、RLHF相比效果如何?

我没有做RLHF实验(感兴趣的读者可以看我之前的文章:https://magazine.sebastianraschka.com/p/llm-training-rlhf-and-its-alternatives ),但我做了全量微调的对比。

全量微调至少需要2张GPU,耗时3.5小时,每张卡显存占用36.66 GB。但最终的基准测试结果并不理想,大概率是过拟合或者超参数设置不佳导致的。

figure16

Q9:多个LoRA权重可以合并吗?

可以,多组LoRA权重是可以合并的。训练过程中,LoRA权重与预训练权重是分开存储的,每次前向传播时再叠加。

在实际应用中,如果你有多组LoRA权重(比如为每个客户单独训练一组),分开存储可以节省磁盘空间。但训练结束后,也可以将LoRA权重与预训练权重合并,得到一个完整的模型。这样在前向传播时就不需要每次都叠加LoRA权重:

weight += (lora_B @ lora_A) * scaling

我们可以按上面的公式完成权重更新,然后保存合并后的权重。

同理,我们也可以叠加多组LoRA权重:

weight += (lora_B_set1 @ lora_A_set1) * scaling_set1
weight += (lora_B_set2 @ lora_A_set2) * scaling_set2
weight += (lora_B_set3 @ lora_A_set3) * scaling_set3
...

我还没有通过实验评估这种方法的效果,但技术上已经可以实现,Lit-GPT提供了对应的脚本:https://github.com/Lightning-AI/lit-gpt/blob/main/scripts/merge_lora.py

Q10:逐层最优秩适配(Layer-wise Optimal Rank Adaptation)效果如何?

为了简化,我们训练深度神经网络时通常给所有层设置相同的学习率,而学习率本身是需要调优的超参数。更进一步,我们也可以给每一层设置不同的学习率(参考:https://kozodoi.me/blog/20220329/discriminative-lr#:~:text=The%20implementation%20of%20layer%2Dwise,with%20the%20corresponding%20learning%20rates. )。但实践中很少这么做,因为这会带来额外开销,而且深度神经网络训练本来就有大量超参数需要调整。

与逐层设置不同学习率类似,我们也可以给不同层设置不同的LoRA秩。我还没有看到相关的实验研究,但有一篇文章详细介绍了这种方法:https://medium.com/@tom_21755/llm-optimization-layer-wise-optimal-rank-adaptation-lora-1444dfbc8e6a(也缩写为LORA)。理论上这听起来是个好主意,但它会让超参数优化的选择数量大幅增加。


本杂志是我的个人兴趣项目。如果希望支持我的创作,欢迎购买我的新书:https://amzn.to/4fqvn0D (我相信这本书会让你收获颇丰,它对大语言模型工作原理的讲解深度是其他地方找不到的。)

figure17

《从零构建大语言模型》购买链接:https://amzn.to/4fqvn0D

如果你读过这本书,能抽出几分钟时间留下评价的话,我会非常感激:https://www.amazon.com/Build-Large-Language-Model-Scratch/dp/1633437167 。这对作者的帮助非常大!

另外,我最近也在Substack上开通了付费订阅选项,可以直接支持本杂志的创作。

《Ahead of AI》是读者支持的出版物。若想接收新文章并支持我的创作,欢迎成为免费或付费订阅者。

Leave a Reply

Your email address will not be published. Required fields are marked *

*