一張RTX 2080Ti搞定大模型訓(xùn)練！算力節(jié)省136倍！

任何 transformer 變體、任何數(shù)據(jù)集都通用。

在一塊消費(fèi)級 GPU 上只用一天時(shí)間訓(xùn)練，可以得到什么樣的 BERT 模型？

最近一段時(shí)間，語言模型再次帶火了 AI 領(lǐng)域。預(yù)訓(xùn)練語言模型的無監(jiān)督訓(xùn)練屬性使其可以在海量樣本基礎(chǔ)上進(jìn)行訓(xùn)練，并獲得大量語義語法知識，不論分類還是問答，似乎沒有 AI 解決不了的問題。

然而，大模型既帶來了技術(shù)突破，也對算力提出了無窮無盡的需求。

最近，來自馬里蘭大學(xué)的 Jonas Geiping、Tom Goldstein 討論了所有關(guān)于擴(kuò)大計(jì)算規(guī)模的研究，深入探討了縮小計(jì)算規(guī)模的改進(jìn)方向。他們的研究引發(fā)了機(jī)器學(xué)習(xí)社區(qū)的關(guān)注。

在新研究中，作者對于單塊消費(fèi)級 GPU（RTX 2080Ti）能訓(xùn)練出什么樣的語言模型進(jìn)行了討論，并獲得了令人興奮的結(jié)果。讓我們看看它是如何實(shí)現(xiàn)的：

模型規(guī)模的擴(kuò)展

在自然語言處理（NLP）領(lǐng)域，基于 Transformer 架構(gòu)的預(yù)訓(xùn)練模型已經(jīng)成為主流，并帶來諸多突破性進(jìn)展。很大程度上，這些模型性能強(qiáng)大的原因是它們的規(guī)模很大。隨著模型參數(shù)量和數(shù)據(jù)量的增長，模型的性能會(huì)不斷提高。因此，NLP 領(lǐng)域內(nèi)掀起了一場增大模型規(guī)模的競賽。

然而，很少有研究人員或從業(yè)者認(rèn)為他們有能力訓(xùn)練大型語言模型（LLM），通常只有行業(yè)內(nèi)的科技巨頭擁有訓(xùn)練 LLM 的資源。

為了扭轉(zhuǎn)這一趨勢，來自馬里蘭大學(xué)的研究者進(jìn)行了一番探索。

論文《Cramming: Training a Language Model on a Single GPU in One Day》：

論文鏈接：https://arxiv.org/abs/2212.14034

這個(gè)問題對于大多數(shù)研究人員和從業(yè)者來說具有重要意義，因?yàn)檫@將成為模型訓(xùn)練成本的參考，并有望打破 LLM 訓(xùn)練成本超高的瓶頸。該研究的論文迅速在推特上引發(fā)關(guān)注和討論。

IBM 的 NLP 研究專家 Leshem Choshen 在推特上評價(jià)道：「這篇論文總結(jié)了所有你能想到的大模型訓(xùn)練 trick。」

馬里蘭大學(xué)的研究者認(rèn)為：如果按比例縮小的模型預(yù)訓(xùn)練是大型預(yù)訓(xùn)練的可行模擬，那么這將開啟一系列目前難以實(shí)現(xiàn)的大規(guī)模模型的進(jìn)一步學(xué)術(shù)研究。

此外，該研究嘗試對過去幾年 NLP 領(lǐng)域的整體進(jìn)展進(jìn)行基準(zhǔn)測試，而不僅僅局限于模型規(guī)模的影響。

該研究創(chuàng)建了一項(xiàng)稱為「Cramming」的挑戰(zhàn) —— 在測試前一天學(xué)習(xí)整個(gè)語言模型。研究者首先分析了訓(xùn)練 pipeline 的方方面面，以了解哪些修改可以實(shí)際提高小規(guī)模模擬模型的性能。并且，該研究表明，即使在這種受限環(huán)境中，模型性能也嚴(yán)格遵循在大型計(jì)算環(huán)境中觀察到的擴(kuò)展定律。

雖然較小的模型架構(gòu)可以加快梯度計(jì)算，但隨著時(shí)間的推移，模型改進(jìn)的總體速度幾乎保持不變。該研究嘗試?yán)脭U(kuò)展定律在不影響模型大小的情況下通過提高梯度計(jì)算的有效率獲得性能提升。最后，該研究成功訓(xùn)練出性能可觀的模型 —— 在 GLUE 任務(wù)上接近甚至超過 BERT—— 而且訓(xùn)練成本很低。

資源有限

為了模擬普通從業(yè)者和研究人員的資源環(huán)境，該研究首先構(gòu)建了一個(gè)資源受限的研究環(huán)境：

一個(gè)任意大小的基于 transformer 的語言模型，完全從頭開始使用掩碼語言建模（masked-language modeling）進(jìn)行訓(xùn)練；

pipeline 中不能包含現(xiàn)有的預(yù)訓(xùn)練模型；

任何原始文本（不包括下游數(shù)據(jù)）都可以包含在訓(xùn)練中，這意味著可以通過明智地選擇如何以及何時(shí)對數(shù)據(jù)進(jìn)行采樣來實(shí)現(xiàn)加速，前提是采樣機(jī)制不需要預(yù)訓(xùn)練模型；

原始數(shù)據(jù)的下載和預(yù)處理不計(jì)入總預(yù)算，這里的預(yù)處理包括基于 CPU 的 tokenizer 構(gòu)造、tokenization 和 filtering，但不包括表征學(xué)習(xí)；

訓(xùn)練僅在單塊 GPU 上進(jìn)行 24 小時(shí)；

下游性能在 GLUE 上進(jìn)行評估，GLUE 上的下游微調(diào)僅限于僅使用下游任務(wù)的訓(xùn)練數(shù)據(jù)進(jìn)行簡單訓(xùn)練（5 個(gè) epoch 或者更少），并且需要使用為所有 GLUE 任務(wù)設(shè)置的全局超參數(shù)，下游微調(diào)不計(jì)算在總預(yù)算中。

該研究與一些經(jīng)典大模型的具體訓(xùn)練設(shè)置比較如下表所示：

改進(jìn)方法

研究人員實(shí)施并測試了已有工作提出的一些修改方向，包括通用實(shí)現(xiàn)和初始數(shù)據(jù)設(shè)置，并嘗試了修改架構(gòu)、訓(xùn)練以及改動(dòng)數(shù)據(jù)集的方法。

實(shí)驗(yàn)在 PyTorch 中進(jìn)行，不使用特質(zhì)化的實(shí)現(xiàn)以盡量公平，所有內(nèi)容都保留在 PyTorch 框架的實(shí)現(xiàn)級別上，只允許可應(yīng)用于所有組件的自動(dòng)運(yùn)算符融合，另外只有在選擇了最終的架構(gòu)變體之后，才會(huì)重新啟用高效注意力內(nèi)核。

圖 1：不同 transformer 架構(gòu)變體 MLM 損失函數(shù)與 token 對比的情況。左：全局視圖。右圖：放大到 10e8 和更多 token 情況下。所有模型都用相同算力成本訓(xùn)練，我們可以看到：通過架構(gòu)重塑實(shí)現(xiàn)的改進(jìn)微乎其微。

有關(guān)提升性能，我們最先想到的方法肯定是修改模型架構(gòu)。從直覺上，較小 / 較低容量的模型似乎在一日一卡式的訓(xùn)練中是最優(yōu)的。然而在研究了模型類型與訓(xùn)練效率之間的關(guān)系后，研究人員發(fā)現(xiàn)縮放法則為縮小規(guī)模設(shè)置了巨大的障礙。每個(gè) token 的訓(xùn)練效率在很大程度上取決于模型大小，而不是 transformer 的類型。

此外，較小的模型學(xué)習(xí)效率較低，這在很大程度上減緩了吞吐量的增加。幸運(yùn)的是，在相同大小的模型中，訓(xùn)練效率幾乎保持不變這一事實(shí)，意味著我們可以在參數(shù)量類似的架構(gòu)中尋找合適的，主要根據(jù)影響單個(gè)梯度步驟的計(jì)算時(shí)間來做出設(shè)計(jì)選擇。

圖 2：學(xué)習(xí)率 Schedule。盡管在全局范圍內(nèi)行為相似，但在中間的放大圖里可以看到差異確實(shí)存在。

在該工作中，作者研究了訓(xùn)練超參數(shù)對 BERT-base 架構(gòu)的影響。可以理解的是，原始 BERT 訓(xùn)練方法的模型在 Cramming 式訓(xùn)練要求中的表現(xiàn)不佳，因此研究人員重新審視了一些標(biāo)準(zhǔn)選擇。

作者也研究了優(yōu)化數(shù)據(jù)集的思路。擴(kuò)展法則阻礙了通過架構(gòu)修改取得重大收益的方式（超出計(jì)算效率），但縮放定律并不妨礙我們在更好的數(shù)據(jù)上進(jìn)行訓(xùn)練。如果想在在每秒訓(xùn)練更多的 token，我們應(yīng)該尋求在更好的 token 上訓(xùn)練。

圖 3：在 bookcorpus-wikipedia 數(shù)據(jù)的 Cramming 式訓(xùn)練中訓(xùn)練的模型的詞匯量與 GLUE 分?jǐn)?shù)和 MNLI 準(zhǔn)確性。

在 GLUE 上的表現(xiàn)

研究人員系統(tǒng)地評估了 GLUE 基準(zhǔn)的性能和 WNLI，并注意到在前面的部分中只使用了 MNLI (m)，并且沒有根據(jù)完整的 GLUE 分?jǐn)?shù)調(diào)整超參數(shù)。在新研究中對于 BERT-base 作者微調(diào)了 5 個(gè) epoch 的所有數(shù)據(jù)集，batch size 為 32，學(xué)習(xí)率為 2 × 10-5。對于 Cramming 訓(xùn)練的模型這不是最優(yōu)的，其可以從 16 的 batch size 和 4 × 10?5 的學(xué)習(xí)率以及余弦衰減中獲得微小的改進(jìn)（此設(shè)置不會(huì)改進(jìn)預(yù)訓(xùn)練的 BERT check point）。

表 3 和表 4 描述了此設(shè)置在 GLUE 下游任務(wù)上的性能。作者比較了原始的 BERT-base check point、在達(dá)到算力上限后停止的 BERT 預(yù)訓(xùn)練設(shè)置、Izsak 等人 2021 年研究中描述的設(shè)置和修改后的設(shè)置，為每塊 GPU 設(shè)置訓(xùn)練一天。總體而言，性能出奇地好，尤其是對于 MNLI、QQP、QNLI 和 SST-2 等較大的數(shù)據(jù)集，下游微調(diào)可以消除完整 BERT 模型和 Cramming 設(shè)置變體之間的剩余差異。

此外，作者發(fā)現(xiàn)新方法與算力有限的普通 BERT 訓(xùn)練及 Izsak 等人描述的方法相比都有很大改進(jìn)。對于 Izsak 等人的研究，其描述的方法最初是為一個(gè)完整的 8 GPU 刀片服務(wù)器設(shè)計(jì)的，并且在新的場景中，將其中的 BERT-large 模型壓縮到較小的 GPU 上是導(dǎo)致大部分性能下降的原因。

表格 3：基線 BERT 與 Cramming 版本模型的 GLUE-dev 性能比較。其中所有任務(wù)的超參數(shù)都是固定的，epoch 限制為 5 個(gè)，缺失值為 NaN。是為 8 GPU 刀片服務(wù)器設(shè)計(jì)的，而在這里，所有計(jì)算被塞進(jìn)了一塊 GPU。

表格 4：基線 BERT 與填充模型的 GLUE-dev 性能比較。

總體而言，使用論文中的方法，訓(xùn)練結(jié)果已經(jīng)非常接近原版 BERT，但要知道后者使用的總 FLOPS 是新方法 45-136 倍（在 16 塊 TPU 上要花費(fèi)四天時(shí)間）。而當(dāng)訓(xùn)練時(shí)間延長 16 倍時(shí)（在 8 塊 GPU 上訓(xùn)練兩天），新方法的性能實(shí)際上比原始 BERT 提高了很多，達(dá)到了 RoBERTa 的水平。

總結(jié)

在該工作中，人們討論了基于 transformer 的語言模型在計(jì)算量非常有限的環(huán)境中可以實(shí)現(xiàn)多少性能，值得慶幸的是，幾條修改方向可以讓我們在 GLUE 上獲得不錯(cuò)的下游性能。研究人員表示，希望這項(xiàng)工作可以為進(jìn)一步的改進(jìn)提供一個(gè)基線，并進(jìn)一步給近年來為 transformer 架構(gòu)提出的許多改進(jìn)和技巧提供理論支撐。

審核編輯 ：李倩