LaVIN-lite:单张消费级显卡微调多模态大模型

把大模型的训练门槛打下来! 我们在单张消费级显卡上实现了多模态大模型(LaVIN-7B, LaVIN-13B)的适配和训练,这篇文章主要介绍一下用到的技术方案和技术细节,供有需要的人参考。这里用到的模型是 LaVIN (语言模型是LLaMA,视觉模型是ViT-L)。LaVIN通过参数高效的训练能将LLaMA拓展到多模态来完成图文问答、对话以及文本对话等等任务。

目前的结果: 7B的多模态大模型训练(LaVIN-7B)大约需要8~9G的显存,13B的多模态大模型训练(LaVIN-13B)大约需要13~14G显存 。目前的模型在单张消费级显卡上已经完全能够完成训练了,性能相较于fp16略有下降,但是仍然极具竞争力!未来预计65B的模型也能在单张A100(40G)上完成训练,我们后续会公布结果。所有的训练和测试代码已经全部开源: github.com/luogen1996/L

技术方案

我们的技术方案结合了LaVIN和qlora,主要分为以下几点:

  • 参数高效的多模态适配 (大概减少了一大半显存)
  • 4bit量化训练 (大概减少了3~8G的固定显存)
  • 梯度累计+gradient checkpointing (大概减少了一半多的显存)
  • Paged Optimizer (作用不是很明显)

参数高效的多模态适配。

在此之前,我先简单介绍一下之前的工作《 Cheap and Quick: Efficient Vision-Language Instruction Tuning for Large Language Models 》。我们在这个工作中提出了一种参数高效的适配方法,能够在将整个LLM参数冻住的情况下实现:

  1. 参数高效的多模态大模型适配(仅花费3~6M额外参数)
  2. 端到端高效训练 (减少2/3的训练时间)
  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来进行适配,也欢迎大家基于这个基线来进行进一步探索和比较

最后,在解决训练的问题之后,我们会持续推进模型能力的提升以及应用场景的创新。另外,多模态对话模型我们也在持续迭代中,未来也会以技术报告的形式来进行分享。

编辑于 2023-07-25 23:20 ・IP 属地广东

文章被以下专栏收录