論文鏈接： https://arxiv.org/abs/2307.05695
代碼倉庫： https://github.com/guitaricet/peft_pretraining

一段時間以來，大模型（LLMs）社區(qū)的研究人員開始關(guān)注于如何降低訓(xùn)練、微調(diào)和推理LLMs所需要的龐大算力，這對于繼續(xù)推動LLMs在更多的垂直領(lǐng)域中發(fā)展和落地具有非常重要的意義。目前這一方向也有很多先驅(qū)工作，例如從模型結(jié)構(gòu)上創(chuàng)新的RWKV，直接替換計算量較大的Transformer架構(gòu)，改用基于RNN范式的新架構(gòu)。還有一些方法從模型微調(diào)階段入手，例如在原有LLMs中加入?yún)?shù)量較小的Adapter模塊來進行微調(diào)。還有微軟提出的低秩自適應(yīng)（Low-Rank Adaptation，LoRA）方法，LoRA假設(shè)模型在任務(wù)適配過程中對模型權(quán)重的更新量可以使用低秩矩陣進行估計，因而可以用來間接優(yōu)化新加入的輕量級適應(yīng)模塊，同時保持原有的預(yù)訓(xùn)練權(quán)重不變。目前LoRA已經(jīng)成為大模型工程師必備的一項微調(diào)技能，但本文作者仍然不滿足于目前LoRA所能達(dá)到的微調(diào)效果，并進一步提出了一種可疊加的低秩微調(diào)方法，稱為ReLoRA。

本文來自馬薩諸塞大學(xué)洛厄爾分校的研究團隊，作者團隊將ReLoRA應(yīng)用在具有高達(dá)350M參數(shù)的Transformer上時，展現(xiàn)出了與常規(guī)神經(jīng)網(wǎng)絡(luò)訓(xùn)練相當(dāng)?shù)男阅?/strong>。此外，本文作者還觀察到ReLoRA的微調(diào)效率會隨著模型參數(shù)規(guī)模的增加而不斷提高，這使得其未來有可能成為訓(xùn)練超大規(guī)模（通常超過1B參數(shù)）LLMs的新型手段。

ReLoRA通過序列疊加的方式僅訓(xùn)練一小組參數(shù)就可以實現(xiàn)與全秩訓(xùn)練相當(dāng)?shù)男阅埽?strong>并且遵循LoRA方法的基礎(chǔ)原則，即保持原始網(wǎng)絡(luò)的凍結(jié)權(quán)重并添加新的可訓(xùn)練參數(shù)。乍一看，這種方式可能顯得計算效率低下，但我們需要清楚的是，這種方法可以通過減小梯度和優(yōu)化器狀態(tài)的大小，來顯著提高顯存效率。例如Adam優(yōu)化器狀態(tài)消耗的顯存通常是模型權(quán)重占用的兩倍。通過大幅減少可訓(xùn)練參數(shù)的數(shù)量，ReLoRA可以在相同的顯存條件下使用更大的batchsize大小，從而最大限度地提高硬件效率，ReLoRA的整體操作細(xì)節(jié)如下圖所示。

03. 實驗效果

為了清晰的評估ReLoRA方法的性能，作者將其應(yīng)用在各種規(guī)模大小（60M、130M、250M 和 350M）的Transformer模型上，并且都在C4數(shù)據(jù)集上進行訓(xùn)練和測試。為了展現(xiàn)ReLoRA方法的普適性，作者重點考察NLP領(lǐng)域的基礎(chǔ)語言建模任務(wù)。模型架構(gòu)和訓(xùn)練超參數(shù)設(shè)置基本與LLaMA模型保持一致。與LLaMA不同的是，作者在實驗中將原始的注意力機制（使用float32進行 softmax計算）替換為了Flash注意力[2]，并且使用bfloat16精度進行計算，這樣操作可以將訓(xùn)練吞吐量提高50-100%，且沒有任何訓(xùn)練穩(wěn)定性問題。此外，使用ReLoRA方法訓(xùn)練的模型參數(shù)規(guī)模相比LLaMA要小得多，最大的模型參數(shù)才僅有350M，使用8個RTX4090上訓(xùn)練了一天時間就可以完成。

下圖展示了本文方法與其他方法的性能對比效果，可以看到ReLoRA顯著優(yōu)于低秩LoRA方法，證明了我們提出的修改的有效性。此外，ReLoRA還實現(xiàn)了與滿秩訓(xùn)練（Full training）相當(dāng)?shù)男阅?，并且我們可以觀察到，隨著網(wǎng)絡(luò)規(guī)模的增加，性能差距逐漸縮小。有趣的是，ReLoRA 唯一無法超過的基線模型是僅具有60M參數(shù)的最小模型。這一觀察結(jié)果表明，ReLoRA在改進大型網(wǎng)絡(luò)的訓(xùn)練方面更加有效，這與作者最開始研究探索一種改進大型網(wǎng)絡(luò)訓(xùn)練方法的目標(biāo)是一致的。

04. 總結(jié)

本文是一篇專注于減輕大型Transformer語言模型訓(xùn)練代價的工作，作者選取了一條非常具有前景的方向，即低秩訓(xùn)練技術(shù)，并且從最樸素的低秩矩陣分解 (LoRA) 方法出發(fā)，利用多個疊加的低秩更新矩陣來訓(xùn)練高秩網(wǎng)絡(luò)，為了實現(xiàn)這一點，作者精心設(shè)計了包含參數(shù)重新啟動、鋸齒狀學(xué)習(xí)率調(diào)度算法和優(yōu)化器參數(shù)重置等一系列操作，這些操作共同提高了ReLoRA算法的訓(xùn)練效率，在某些情況下甚至能夠達(dá)到與全秩訓(xùn)練相當(dāng)?shù)男阅?/strong>，尤其實在超大規(guī)模的Transformer網(wǎng)絡(luò)中。作者通過大量的實驗證明了ReLoRA的算法可行性和操作有效性，不知ReLoRA是否也會成為大模型工程師一項必備的算法技能呢？