把大模型的訓(xùn)練門(mén)檻打下來(lái)！我們?cè)趩螐埾M(fèi)級(jí)顯卡上實(shí)現(xiàn)了多模態(tài)大模型（LaVIN-7B, LaVIN-13B）的適配和訓(xùn)練，這篇文章主要介紹一下用到的技術(shù)方案和技術(shù)細(xì)節(jié)，供有需要的人參考。這里用到的模型是LaVIN(語(yǔ)言模型是LLaMA，視覺(jué)模型是ViT-L)。LaVIN通過(guò)參數(shù)高效的訓(xùn)練能將LLaMA拓展到多模態(tài)來(lái)完成圖文問(wèn)答、對(duì)話以及文本對(duì)話等等任務(wù)。

目前的結(jié)果：7B的多模態(tài)大模型訓(xùn)練（LaVIN-7B）大約需要8~9G的顯存，13B的多模態(tài)大模型訓(xùn)練（LaVIN-13B）大約需要13~14G顯存。目前的模型在單張消費(fèi)級(jí)顯卡上已經(jīng)完全能夠完成訓(xùn)練了，性能相較于fp16略有下降，但是仍然極具競(jìng)爭(zhēng)力！未來(lái)預(yù)計(jì)65B的模型也能在單張A100（40G）上完成訓(xùn)練，我們后續(xù)會(huì)公布結(jié)果。

技術(shù)方案

我們的技術(shù)方案結(jié)合了LaVIN和qlora，主要分為以下幾點(diǎn)：

參數(shù)高效的多模態(tài)適配 (大概減少了一大半顯存)

4bit量化訓(xùn)練 （大概減少了3~8G的固定顯存）

梯度累計(jì)+gradient checkpointing （大概減少了一半多的顯存）

Paged Optimizer (作用不是很明顯)

參數(shù)高效的多模態(tài)適配。

在此之前，我先簡(jiǎn)單介紹一下之前的工作《Cheap and Quick: Efficient Vision-Language Instruction Tuning for Large Language Models》。我們?cè)谶@個(gè)工作中提出了一種參數(shù)高效的適配方法，能夠在將整個(gè)LLM參數(shù)凍住的情況下實(shí)現(xiàn)：

參數(shù)高效的多模態(tài)大模型適配（僅花費(fèi)3~6M額外參數(shù)）

端到端高效訓(xùn)練 （減少2/3的訓(xùn)練時(shí)間）

單模態(tài)和多模態(tài)的自動(dòng)切換（兼容不同模態(tài)）

通過(guò)這種方式，我們?cè)赟cienceQA上達(dá)到了接近SOTA的性能，同時(shí)實(shí)現(xiàn)了文本模態(tài)和圖文模態(tài)的同時(shí)適配。這種參數(shù)高效的訓(xùn)練方式，實(shí)際上節(jié)約了大部分的顯存。以LLaVA為比較對(duì)象，在完全微調(diào)大模型的情況下，LLaVA-13B在A100（80G）上會(huì)爆顯存。相比之下，LaVIN-13B僅僅需要大約55G的顯存開(kāi)銷。考慮到LLaVA還使用了gradient checkpointing，LaVIN-13B至少節(jié)省了一半的顯存開(kāi)銷（估計(jì)），同時(shí)訓(xùn)練速度會(huì)更快。相比于現(xiàn)有的參數(shù)高效的方法，我們的方案在性能和適配性上有顯著優(yōu)勢(shì)，具體參考論文，這里不贅述了。但是由于deepspeed好像不支持參數(shù)高效的訓(xùn)練方式，所以實(shí)際中顯存開(kāi)銷其實(shí)和加滿優(yōu)化的LLaVA差不多,甚至略多一點(diǎn)。

4bit量化訓(xùn)練

4bit量化訓(xùn)練主要參考了qlora。簡(jiǎn)單來(lái)說(shuō)，qlora把LLM的權(quán)重量化成了4bit來(lái)存儲(chǔ)，同時(shí)在訓(xùn)練過(guò)程中反量化成16bit來(lái)保證訓(xùn)練精度。通過(guò)這種方式，能夠大大降低訓(xùn)練過(guò)程中的顯存開(kāi)銷（訓(xùn)練速度應(yīng)該區(qū)別不大）。這種方法非常適合和參數(shù)高效的方法進(jìn)行結(jié)合。但是原文中針對(duì)的是單模態(tài)的LLM，同時(shí)代碼已經(jīng)封在了huggingface的庫(kù)里。因此，我們把核心代碼從huggingface的庫(kù)里抽取出來(lái)，同時(shí)遷移到了LaVIN的代碼里。主要原理就是將LLM中所有的線性層替換成4bit的量化層，感興趣的可以去參考一下我們的寫(xiě)法，在quantization.py以及mm_adaptation.py中大概十來(lái)行代碼。

4bit量化訓(xùn)練之后，顯存在bs>1的時(shí)候下降的不是特別明顯。LaVIN-7B大概下降了4~6G的樣子，但是這部分的顯存下降是固定的，其實(shí)非常有價(jià)值。到這里我也很好奇qlora怎么把模型塞到單卡里的，這個(gè)時(shí)候LaVIN-7B的顯存開(kāi)銷大概還在36+G的水平。后面check了一下他們的代碼發(fā)現(xiàn)了接下來(lái)的關(guān)鍵設(shè)置。

梯度累計(jì)+gradient checkpointing

這里的關(guān)鍵就在于時(shí)間換空間。通過(guò)batch size (bs)=1+梯度累計(jì)以及gradient checkpointing的方式能夠大大降低顯存開(kāi)銷。這也是qlora訓(xùn)練時(shí)的一大核心（其實(shí)光靠量化訓(xùn)練很難做到顯存的極致壓縮）。我們的實(shí)驗(yàn)結(jié)果大概是這樣：LaVIN-7B在bs=4改成batch size (bs)=1+梯度累計(jì)之后顯存降低到了25G左右。經(jīng)過(guò)gradient checkpointing，顯存降低到9~10G左右。到這里，顯存從原來(lái)的上百G壓縮到了10G左右，已經(jīng)非常可觀了。但是這一步的代價(jià)是訓(xùn)練速度明顯變慢了，但其實(shí)和qlora原文中的速度下降比例差不多。相比于原來(lái)完全訓(xùn)不了的情況來(lái)說(shuō)，這些額外的時(shí)間開(kāi)銷顯得非常微不足道。

Paged Optimizer

Paged Optimizer的作用是在快爆顯存的時(shí)候，會(huì)將optimizer中的一部分權(quán)重遷移到cpu上，從而保證訓(xùn)練的正常進(jìn)行。實(shí)際使用中，沒(méi)有感覺(jué)到太大的區(qū)別。我猜測(cè)是在顯存開(kāi)銷和顯卡顯存非常接近的時(shí)候，這個(gè)設(shè)置能救下急。正常情況下，好像沒(méi)有什么太大的幫助。感興趣的可以試一下8 bit的optimizer，或許幫助更明顯。

性能比較

ScienceQA（多模態(tài)科學(xué)問(wèn)答數(shù)據(jù)集）：在ScienceQA上，我們單卡的情況下完成了4bit訓(xùn)練并和16bit的方法進(jìn)行了比較，結(jié)果如下：

可以看到LaVIN-lite性能仍然遠(yuǎn)超參數(shù)高效的方法LLaMA-Adapter，但是相比較16bit訓(xùn)練的LaVIN，性能出現(xiàn)了略微的下降。我們猜測(cè)原因是4bit訓(xùn)練的時(shí)候可能需要插入更多的adapter來(lái)進(jìn)行適配，也歡迎大家基于這個(gè)基線來(lái)進(jìn)行進(jìn)一步探索和比較。

最后，在解決訓(xùn)練的問(wèn)題之后，我們會(huì)持續(xù)推進(jìn)模型能力的提升以及應(yīng)用場(chǎng)景的創(chuàng)新。另外，多模態(tài)對(duì)話模型我們也在持續(xù)迭代中，未來(lái)也會(huì)以技術(shù)報(bào)告的形式來(lái)進(jìn)行分享。

審核編輯：劉清