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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

資源有限

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

• 一個(gè)任意大小的基于 transformer 的語(yǔ)言模型，完全從頭開(kāi)始使用掩碼語(yǔ)言建模（masked-language modeling）進(jìn)行訓(xùn)練；
• pipeline 中不能包含現(xiàn)有的預(yù)訓(xùn)練模型；
• 任何原始文本（不包括下游數(shù)據(jù)）都可以包含在訓(xùn)練中，這意味著可以通過(guò)明智地選擇如何以及何時(shí)對(duì)數(shù)據(jù)進(jìn)行采樣來(lái)實(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)行評(píng)估，GLUE 上的下游微調(diào)僅限于僅使用下游任務(wù)的訓(xùn)練數(shù)據(jù)進(jìn)行簡(jiǎn)單訓(xùn)練（5 個(gè) epoch 或者更少），并且需要使用為所有 GLUE 任務(wù)設(shè)置的全局超參數(shù)，下游微調(diào)不計(jì)算在總預(yù)算中。

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

改進(jìn)方法

研究人員實(shí)施并測(cè)試了已有工作提出的一些修改方向，包括通用實(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)級(jí)別上，只允許可應(yīng)用于所有組件的自動(dòng)運(yùn)算符融合，另外只有在選擇了最終的架構(gòu)變體之后，才會(huì)重新啟用高效注意力內(nèi)核。

圖 1：不同 transformer 架構(gòu)變體 MLM 損失函數(shù)與 token 對(duì)比的情況

左：全局視圖。右圖：放大到 10e8 和更多 token 情況下。

所有模型都用相同算力成本訓(xùn)練，我們可以看到：通過(guò)架構(gòu)重塑實(shí)現(xiàn)的改進(jìn)微乎其微。

有關(guān)提升性能，我們最先想到的方法肯定是修改模型架構(gòu)。從直覺(jué)上，較小 / 較低容量的模型似乎在一日一卡式的訓(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í)間來(lái)做出設(shè)計(jì)選擇。

圖 2：學(xué)習(xí)率 Schedule

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

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

作者也研究了優(yōu)化數(shù)據(jù)集的思路。擴(kuò)展法則阻礙了通過(guò)架構(gòu)修改取得重大收益的方式（超出計(jì)算效率），但縮放定律并不妨礙我們?cè)诟玫臄?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)地評(píng)估了 GLUE 基準(zhǔn)的性能和 WNLI，并注意到在前面的部分中只使用了 MNLI (m)，并且沒(méi)有根據(jù)完整的 GLUE 分?jǐn)?shù)調(diào)整超參數(shù)。在新研究中對(duì)于 BERT-base 作者微調(diào)了 5 個(gè) epoch 的所有數(shù)據(jù)集，batch size 為 32，學(xué)習(xí)率為 2 × 10-5。對(duì)于 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)練一天。總體而言，性能出奇地好，尤其是對(duì)于 MNLI、QQP、QNLI 和 SST-2 等較大的數(shù)據(jù)集，下游微調(diào)可以消除完整 BERT 模型和 Cramming 設(shè)置變體之間的剩余差異。 

此外，作者發(fā)現(xiàn)新方法與算力有限的普通 BERT 訓(xùn)練及 Izsak 等人描述的方法相比都有很大改進(jìn)。對(duì)于 Izsak 等人的研究，其描述的方法最初是為一個(gè)完整的 8 GPU 刀片服務(wù)器設(shè)計(jì)的，并且在新的場(chǎng)景中，將其中的 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í)間延長(zhǎng) 16 倍時(shí)（在 8 塊 GPU 上訓(xùn)練兩天），新方法的性能實(shí)際上比原始 BERT 提高了很多，達(dá)到了 RoBERTa 的水平。

總結(jié)

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