LongQLoRA：单卡高效扩展LLaMA2-13B的上下文长度
本文将介绍我们关于低资源高效扩展大模型上下文长度的工作：LongQLoRA。其中将涉及Position Interpolation与QLoRA的相关知识，建议大家结合往期文章，有助于理解本工作：
论文链接：
https://arxiv.org/abs/2311.04879
项目地址：
https://github.com/yangjianxin1/LongQLoRA
01
论文摘要
扩展大模型的上下文长度是一项非常重要的工作，在长文档问答、长文本摘要和RAG等多种应用场景中，大模型的上下文长度一度成为瓶颈。
我们提出了LongQLoRA，这是一种低资源高效的扩展大模型上下文长度的方法。LongQLoRA结合了Position Interpolation、QLoRA与Shift Short Attention的优点，仅需一张32GB的V100 GPU，就可以将LLaMA2-7B和13B模型的上下文长度从4096扩展到8192，甚至是12k，并且仅需微调1000个step。
LongQLoRA在PG19和 Proof-pile数据集上的困惑度表现具有很强的竞争力。在上述评测集中，我们的模型优于LongLoRA，在8192的评测长度的表现及其接近MPT-7B-8K。
我们构建了一个包含39k长指令的数据集，并将Vicuna-13B的上下文长度从4096扩展到8192，模型在Long Context与Short Context的生成任务中均取得良好的表现。
我们也进行了一些消融实验，包括LoRA Rank、微调步数、推理注意力机制等，以找出更优的训练配置。
02
Introduction
随着LLaMA系列模型的出现，涌现出了一系列基于LLaMA的开源工作，例如Vicuna、Gunaco、WizardLM等。经过指令微调与RLHF，这些开源工作在许多任务上都取得了优异的表现，并且在部分任务中能够媲美甚至超越ChatGPT。
LLaMA系列模型在预定的上下文长度中进行训练，例如LLaMA为2048，LLaMA2为4096。由于LLaMA系列模型采用RoPE旋转位置编码，所以具有较弱的长度外推性，当模型的输入长度超过预定长度时，模型的困惑度将会急剧提升，性能急剧下降。对于一些长文本任务，例如多文档问答、长文本摘要、长历史对话等，有限的上下文长度限制了模型的应用场景。
为了提高模型的上下文长度，最直接的方法就是使用更长的输入文本对模型进行训练。然而，这种方法需要消耗非常多的训练资源，并且收敛速度较慢，训练效果不佳。为了解决该问题，Meta提出了Position Interpolation，使用32 A100 GPU将LLaMA的上下文长度从2048扩展至8192，且仅需微调1000步即可获良好的效果。Focused Transformer (FOT)开源了LongLLaMA，将模型长度扩展至256k。FOT是一种即插即用的扩展方法，模型很容易外推到更长的序列，例如在8K token上训练的模型，可以很容易外推到256K窗口大小。LongLLaMA也使用了128 TPU进行训练。LongLoRA提出了Shift Short Attention，并且结合Position Interpolation和LoRA，实现了一种更高效的长度扩展方法，他们在8 A100上，将LLaMA模型的长度扩展至100k。
然而上述方法中的Position Interpolation、FOT都需要消耗大量的训练资源，相比之下LongLoRA能够节省很多训练资源，但对于大部分普通研究者仍然难以承担。是否有其他方法能够进一步节省训练资源，且保证训练效果呢？
我们认为QLoRA是一个很好的选择，QLoRA是一种高效的微调方法，它将预训练模型的权重量化为4-bit，以降低模型的显存占用，然后插入可学习的LoRA adapters。该方法能够在48GB GPU上对65B模型进行微调，并且能够取得与全量参数微调非常接近的效果。
在本文中，我们提出了LongQLoRA，这是一种低资源高效的扩展大模型上下文长度的方法，我们仅使用一张V100，即可将LLaMA2-7B的上下文长度从4096扩展至8192。LongQLoRA主要结合了QLoRA、Position Interpolation与Shift Short Attention等方法的优势。在训练细节上，LongQLoRA与LongLoRA存在差异。在LongQLoRA中，所有layer均插入adapter，增大LoRA Rank至64，且不训练embeddings和normalization层。
除此之外，我们也构造了一个包含39k条长指令的数据集，主要包含Book Sum, Alpaca, WizardLM, Natural Questions等指令数据。我们使用该数据将Vicuna-13B的上下文长度扩展至8196，并且取得了不错的效果。
我们的工作总结如下：
1. 提出LongQLoRA，将QLoRA、Position Interpolation和 Shift Short Attention相结合，更加节省训练资源，并且取得了优秀表现。仅使用单张V100，即可将LLaMA2-7B和Vicuna-13B扩展至8192。
2. 我们在PG19和Proof-pile数据集进行评测，证明了LongQLoRA的有效性，其表现优于LongLoRA，且逼近MPT-7B-8K。
3. 我们收集和构造了5.4k长文本预训练数据和3.9k长指令数据，并且将我们的数据、训练代码、模型权重的开源。
03
方法介绍
Background
Position Interpolation是由Meta提出的一种扩展大模型上下文长度的方法，经过1000步长文本微调即可将LLaMA的上下文长度扩展至32768。Position Interpolation通过缩小RoPE旋转角度，以达到上下文扩展的目的。
经过位置线性插值之后，仅需要使用少量长文本数据对模型进行1000步左右的微调，即可在PG19数据集上取得优秀的困惑度表现。
Shift Short Attention是由LongLoRA提出的一种局部注意力机制，与全局注意力相比，它能够大幅减少显存占用。Shift Short Attention将输入划分为多个group，selft-attention仅作用于每个group内部，假设有
为了弥补局部注意力的不足，Shift Short Attention还会在group之间进行attention计算，实验表明这一优化，能够取得与全局注意力相近的表现。
QLoRA是由华盛顿大学提出的一种高效微调大模型的方法，可在单张A100上对LLaMA-65B进行微调。使用QLoRA微调的LLaMA-65B，可达到ChatGPT性能水平的99.3%（由GPT-4进行评价），并且实验表明QLoRA的性能可以逼近全量参数微调。
QLoRA将预训练模型的权重量化为4-bit以节省显存，然后插入可学习的adapter对模型进行训练。QLoRA提出了4-bit NormalFloat、Double Quantization和Paged Optimizers。
4-bit NormalFloat：一种理论最优的4-bit的量化数据类型，优于当前普遍使用的FP4与Int4。
Double Quantization：相比于当前的模型量化方法，更加节省显存空间。每个参数平均节省0.37bit，对于LLaMA-65B，大约能节省3GB显存空间。
Paged Optimizers：使用NVIDIA统一内存来避免在处理小批量的长序列时出现的梯度检查点内存峰值。
基于上述优化，QLoRA能够在非常少的训练资源的情况下，对大模型进行微调，并且能够取得接近全量参数finetune的效果。
LongQLoRA
LongQLoRA结合了Position Interpolation、QLoRA与Shift Short Attention等方法的优势。我们首先使用Position Interpolation将LLaMA2的上下文长度扩展至目标长度。在训练时，首先使用QLoRA技术将预训练模型的权重量化为4-bit以节省显存。为了进一步节省训练中的显存占用，我们采用了Shift Short Attention。
为了弥补量化所带来的的精度损失，我们在所有layer中均插入adapter，增大LoRA Rank至64。与LongLoRA不同的是，我们并不需要训练embeddings和normalization layer，也可获得很好的效果，这与LoRA Layer与LoRA Rank的设置有关。
在推理时，我们发现采用全局注意力比Shift Short Attention效果更好，所以在推理阶段统一采用全局注意力。这也意味在推理阶段，我们的模型可以与现有的技术进行无缝兼容，例如Flash Attention、vLLM等，无需额外的适配成本。
04
实验设置 & 数据集
我们主要在7B和13B模型上进行实验，在整个实验过程中只使用了一张32GB显存的V100 GPU。我们将LLaMA2-7B和Vicuna-13B的上下文长度从4096扩展到8192。
我们首先使用Position Interpolation技术，缩小RoPE的旋转角度，将模型的上下文长度从4096扩展至8192。
对于QLoRA而言，我们将预训练模型的权重量化为4-bit NormalFloat，LoRA Rank设为64，并且在所有layer都插入LoRA adapter，最终7B与13B模型的训练参数量分别约为1.5M和2.5M。使用Paged Optimizers，对于7B和13B模型，学习率分别设为2e-4和1e-4, warmup step为20，batch size为1，gradient accumulation steps为16。我们采用Deepspeed Zero2策略。对于LLaMA2-7B，我们训练1000个step，而对于Vicuna-13B，我们训练一个epoch。
在微调LLaMA2-7B时，采用next token prediction task，而在微调Vicuna-13B时，仅计算target部分的loss。
对于Shift Short Attention而言，我们将group size设为目标长度的1/4，且仅在训练时使用Shift Short Attention，推理时使用标准的全局注意力机制。
对LLaMA2-7B进行上下文长度扩展时，我们从Redpajama数据集中采样了大约54k条数据，token长度从4096至32768不等。
我们在PG19和Proof-pile数据集上评测模型的困惑度。当评测长度为8192时，我们将模型权重量化为4-bit，对于其他评测长度则使用16-bit进行推理。在评测时，滑动窗口的步长与评测长度相同。
除此之外，我们还构建了一个Long Context的指令数据集，用于对chat模型进行有监督微调。该数据集包含39k条指令数据，主要包含Book Sum、LongQA、WizardLM、Natural Questions等数据，为了适配8192的目标长度，该数据集中的token长度最大均为8192。数据分布如下图所示：
05
实验结果
主要结果
模型在PG19验证集上的困惑度表现如下表所示。当评测长度为8192时，LongQLoRA优于LongLoRA-LoRA和MPT-7B-8K，并且非常接近LongLoRA-Full，困惑度仅高0.03。评测长度为1024到4096时，LongQLoRA也略优于LongLoRA-LoRA，并且接近LongLoRA-Full。
在下表中，展示了模型在Proof-pile测试集上的评测结果。LongQLoRA在1024到8192的上下文长度中均略优于LongLoRA-LoRA，并且也非常接近LongLoRA-Full，在8192的评测长度中几乎与LongLoRA-Full相同。
值得注意的是，MPT-7B-8K是由MPT-7B增量预训练而来，将序列长度扩展为8K，并使用了500B token进行全量参数增量预训练，因此训练成本非常高。LongLoRA虽然节省了许多训练资源，但也消耗了8个A100 GPU。LongQLoRA则更为高效，仅在一张V100上微调1000步，即可获得接近甚至超越其他模型的效果。
LongQLoRA-Vicuna-13B-8K是在Vicuna-13B的基础上进行微调的，其上下文长度从4096扩展到8192。该模型在Long Context生成中也取得了良好的性能，并且避免了短指令遵循能力的退化。下图展示了模型的生成样例。
消融实验
为了研究LoRA Rank的影响，我们使用不同的LoRA Rank对LLaMA2-7B微调1000步。如下表所示，随着LoRA Rank的增大，困惑度逐渐下降。当LoRA Rank设为64时，LongQLoRA实现了与LongLoRA-Full几乎相同的性能，甚至略优于MPT-7B-8K。因此，64是LongQLoRA中一个较为合适的LoRA Rank。也许随着LoRA Rank的进一步增大，困惑度会进一步下降。
我们也对微调步数进行了消融实验。固定评测长度为8192，微调步数从0增加到1000时，模型在PG19验证集上的困惑度表现如下表所示。当微调步数为0时，困惑度很高，仅经过100步微调，困惑度就显著下降。随着微调步数进一步增加，困惑度不断下降，在大约1000步时，模型基本收敛，并且取得了良好的性能。该消融实验进一步验证了LongQLoRA的高效性，仅需使用少量长文本微调1000步即可获得优秀的性能。
我们对推理阶段采取不同的注意力机制进行了消融实验。如下表所示，Shift Short Attention不适合推理阶段，标准全局注意力机制取得更优的表现。这也表明训练时采用Shift Short Attention，而推理时采用标准全局注意力机制是可行的，并且可以完美兼容现有的推理技术和框架。
06
结语
我们提出了LongQLoRA，这是一种低资源高效的扩展大模型上下文长度的方法。我们在LLaMA2-7B与Vicuna-13B上验证了该方法的有效性，并且该方法可以无缝兼容现有的推理技术与框架。
由于训练资源的限制，我们仅在8192的长度上进行了实验，没有进一步探索更大的上下文长度。若后续训练资源允许，我们将尝试在更大的上下文长度中探索方法的可行性。
在「知乎」也能找到我们了，搜索「红雨瓢泼」即可。
给公众号添加【星标⭐️】不迷路