LaVIN-lite：单张消费级显卡微调多模态大模型
进NLP群—>加入NLP交流群
Code链接：https://github.com/luogen1996/LaVIN
把大模型的训练门槛打下来！我们在单张消费级显卡上实现了多模态大模型（LaVIN-7B, LaVIN-13B）的适配和训练，这篇文章主要介绍一下用到的技术方案和技术细节，供有需要的人参考。这里用到的模型是LaVIN(语言模型是LLaMA，视觉模型是ViT-L)。LaVIN通过参数高效的训练能将LLaMA拓展到多模态来完成图文问答、对话以及文本对话等等任务。
目前的结果：7B的多模态大模型训练（LaVIN-7B）大约需要8~9G的显存，13B的多模态大模型训练（LaVIN-13B）大约需要13~14G显存。目前的模型在单张消费级显卡上已经完全能够完成训练了，性能相较于fp16略有下降，但是仍然极具竞争力！未来预计65B的模型也能在单张A100（40G）上完成训练，我们后续会公布结果。所有的训练和测试代码已经全部开源：https://github.com/luogen1996/LaVIN
技术方案
我们的技术方案结合了LaVIN和qlora，主要分为以下几点：
参数高效的多模态适配 (大概减少了一大半显存)
4bit量化训练 （大概减少了3~8G的固定显存）
梯度累计+gradient checkpointing （大概减少了一半多的显存）
Paged Optimizer (作用不是很明显)
参数高效的多模态适配。
在此之前，我先简单介绍一下之前的工作《Cheap and Quick: Efficient Vision-Language Instruction Tuning for Large Language Models》。我们在这个工作中提出了一种参数高效的适配方法，能够在将整个LLM参数冻住的情况下实现：
参数高效的多模态大模型适配（仅花费3~6M额外参数）
端到端高效训练 （减少2/3的训练时间）
单模态和多模态的自动切换（兼容不同模态）
通过这种方式，我们在ScienceQA上达到了接近SOTA的性能，同时实现了文本模态和图文模态的同时适配。这种参数高效的训练方式，实际上节约了大部分的显存。以LLaVA为比较对象，在完全微调大模型的情况下，LLaVA-13B在A100（80G）上会爆显存。相比之下，LaVIN-13B仅仅需要大约55G的显存开销。考虑到LLaVA还使用了gradient checkpointing，LaVIN-13B至少节省了一半的显存开销（估计），同时训练速度会更快。 相比于现有的参数高效的方法，我们的方案在性能和适配性上有显著优势，具体参考论文，这里不赘述了。但是由于deepspeed好像不支持参数高效的训练方式，所以实际中显存开销其实和加满优化的LLaVA差不多,甚至略多一点。
4bit量化训练
4bit量化训练主要参考了qlora。简单来说，qlora把LLM的权重量化成了4bit来存储，同时在训练过程中反量化成16bit来保证训练精度。通过这种方式，能够大大降低训练过程中的显存开销（训练速度应该区别不大）。这种方法非常适合和参数高效的方法进行结合。但是原文中针对的是单模态的LLM，同时代码已经封在了huggingface的库里。因此，我们把核心代码从huggingface的库里抽取出来，同时迁移到了LaVIN的代码里。主要原理就是将LLM中所有的线性层替换成4bit的量化层，感兴趣的可以去参考一下我们的写法，在quantization.py以及mm_adaptation.py中大概十来行代码。
4bit量化训练之后，显存在bs>1的时候下降的不是特别明显。LaVIN-7B大概下降了4~6G的样子，但是这部分的显存下降是固定的，其实非常有价值。 到这里我也很好奇qlora怎么把模型塞到单卡里的，这个时候LaVIN-7B的显存开销大概还在36+G的水平。后面check了一下他们的代码发现了接下来的关键设置。
梯度累计+gradient checkpointing
这里的关键就在于时间换空间。通过batch size (bs)=1+梯度累计以及gradient checkpointing的方式能够大大降低显存开销。这也是qlora训练时的一大核心（其实光靠量化训练很难做到显存的极致压缩）。我们的实验结果大概是这样：LaVIN-7B在bs=4改成batch size (bs)=1+梯度累计之后显存降低到了25G左右。经过gradient checkpointing，显存降低到9~10G左右。到这里，显存从原来的上百G压缩到了10G左右，已经非常可观了。但是这一步的代价是训练速度明显变慢了，但其实和qlora原文中的速度下降比例差不多。相比于原来完全训不了的情况来说，这些额外的时间开销显得非常微不足道。
Paged Optimizer
Paged Optimizer的作用是在快爆显存的时候，会将optimizer中的一部分权重迁移到cpu上，从而保证训练的正常进行。实际使用中，没有感觉到太大的区别。我猜测是在显存开销和显卡显存非常接近的时候，这个设置能救下急。正常情况下，好像没有什么太大的帮助。感兴趣的可以试一下8 bit的optimizer，或许帮助更明显。
性能比较
ScienceQA（多模态科学问答数据集）：在ScienceQA上，我们单卡的情况下完成了4bit训练并和16bit的方法进行了比较，结果如下：
可以看到LaVIN-lite性能仍然远超参数高效的方法LLaMA-Adapter，但是相比较16bit训练的LaVIN，性能出现了略微的下降。我们猜测原因是4bit训练的时候可能需要插入更多的adapter来进行适配，也欢迎大家基于这个基线来进行进一步探索和比较。
最后，在解决训练的问题之后，我们会持续推进模型能力的提升以及应用场景的创新。另外，多模态对话模型我们也在持续迭代中，未来也会以技术报告的形式来进行分享。
进NLP群—>加入NLP交流群
微信扫码关注该文公众号作者