一日一卡挑戰：RTX2080Ti搞定大模型訓練，算力省136倍，學界歡呼

在一塊消費級 GPU 上只用一天時間訓練，可以得到什麼樣的 BERT 模型？

最近一段時間，語言模型再次帶火了 AI 領域。預訓練語言模型的無監督訓練屬性使其可以在海量樣本基礎上進行訓練，並獲得大量語義語法知識，不論分類還是問答，似乎沒有 AI 解決不了的問題。

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

最近，來自馬裡蘭大學的 Jonas Geiping、Tom Goldstein 討論了所有關於擴大計算規模的研究，深入探討了縮小計算規模的改進方向。他們的研究引發了機器學習社群的關注。

在新研究中，作者對於單塊消費級 GPU（RTX 2080Ti）能訓練出什麼樣的語言模型進行了討論，並獲得了令人興奮的結果。讓我們看看它是如何實現的：

模型規模的擴展

在自然語言處理（NLP）領域，基於Transformer 架構的預訓練模型已經成為主流，並帶來許多突破性進展。在很大程度上，這些模型性能強大的原因是它們的規模很大。隨著模型參數量和資料量的成長，模型的效能會不斷提高。因此，NLP 領域內掀起了一場增大模型規模的競賽。

然而，很少有研究人員或從業者認為他們有能力訓練大型語言模型（LLM），通常只有業界的科技巨頭擁有訓練 LLM 的資源。

為了扭轉這一趨勢，來自馬裡蘭大學的研究者進行了一番探索。

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

論文連結：https://arxiv.org/abs /2212.14034

這個問題對於大多數研究人員和從業者來說具有重要意義，因為這將成為模型訓練成本的參考，並有望打破LLM 訓練成本超高的瓶頸。研究的論文迅速在推特上引發關注和討論。

IBM 的NLP 研究專家Leshem Choshen 在推特上評論：「這篇論文總結了所有你能想到的大模型訓練trick。」

馬裡蘭大學的研究者認為：如果按比例縮小的模型預訓練是大型預訓練的可行模擬，那麼這將開啟一系列目前難以實現的大規模模型的進一步學術研究。

此外，研究嘗試對過去幾年 NLP 領域的整體進展進行基準測試，而不僅限於模型規模的影響。

該研究創建了一項稱為「Cramming」的挑戰 —— 在測試前一天學習整個語言模型。研究者首先分析了訓練 pipeline 的方方面面，以了解哪些修改可以實際提高小規模模擬模型的效能。並且，該研究表明，即使在這種受限環境中，模型效能也嚴格遵循在大型運算環境中觀察到的擴展定律。

雖然較小的模型架構可以加快梯度運算，但隨著時間的推移，模​​型改進的整體速度幾乎保持不變。該研究嘗試利用擴展定律在不影響模型大小的情況下透過提高梯度計算的有效率來獲得效能提升。最後，研究成功訓練出表現可觀的模型 —— 在 GLUE 任務上接近甚至超過 BERT—— 而且訓練成本很低。

資源有限

為了模擬一般實務工作者和研究人員的資源環境，研究首先建構了一個資源受限的研究環境：

• 一個任意大小的基於transformer 的語言模型，完全從頭開始使用掩碼語言建模（masked-language modeling）進行訓練；
• pipeline 中不能包含現有的預訓練模型;
• 任何原始文字（不包括下游資料）都可以包含在訓練中，這意味著可以透過明智地選擇如何以及何時對資料進行採樣來實現加速，前提是採樣機制不需要預訓練模型；
• 原始資料的下載和預處理不計入總預算，這裡的預處理包括基於CPU 的tokenizer 構造、tokenization 和filtering，但不包括表徵學習；
• 訓練僅在單一區塊GPU 上進行24 小時；
• 下游效能在GLUE 上進行評估，GLUE 上的下游微調僅限於僅使用下游任務的訓練資料進行簡單訓練（5 個epoch 或更少），並且需要使用為所有GLUE 任務設定的全域超參數，下游微調不計算在總預算中。

本研究與一些經典大模型的具體訓練設定比較如下表所示：

改進方法

研究人員實作並測試了一些已有工作提出的修改方向，包括通用實作和初始資料設置，並嘗試了修改架構、訓練以及改變資料集的方法。

實驗在PyTorch 中進行，不使用特質化的實現以盡量公平，所有內容都保留在PyTorch 框架的實現級別上，只允許可應用於所有組件的自動運算符融合，另外只有在選擇了最終的架構變體之後，才會重新啟用高效注意力核心。

圖 1：不同 transformer 架構變體 MLM 損失函數與 token 比較的情況

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

所有模型都以相同算力成本訓練，我們可以看到：透過架構重塑實現的改進微乎其微。

有關提升效能，我們最先想到的方法肯定是修改模型架構。從直覺上，較小 / 較低容量的模型似乎在一日一卡的訓練中是最優的。然而在研究了模型類型與訓練效率之間的關係後，研究人員發現縮放法則為縮小規模設定了巨大的障礙。每個 token 的訓練效率在很大程度上取決於模型大小，而不是 transformer 的類型。

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

圖 2：學習率 Schedule

儘管在全域範圍內行為相似，但在中間的放大圖裡可以看到差異確實存在。

在這個工作中，作者研究了訓練超參數對 BERT-base 架構的影響。可以理解的是，原始 BERT 訓練方法的模型在 Cramming 式訓練要求中的表現不佳，因此研究人員重新審視了一些標準選擇。

作者也研究了優化資料集的想法。擴展法則阻礙了透過架構修改取得重大收益的方式（超出運算效率），但縮放定律並不妨礙我們在更好的資料上進行訓練。如果想在每秒訓練更多的 token，我們應該尋求在更好的 token 上訓練。

圖 3：在 bookcorpus-wikipedia 資料的 Cramming 式訓練中訓練的模型的詞彙量與 GLUE 分數和 MNLI 準確性。

在GLUE 上的表現

研究人員系統地評估了GLUE 基準的性能和WNLI，並注意到在前面的部分中只使用了MNLI (m)，並且沒有根據完整的GLUE 分數調整超參數。在新研究中對於 BERT-base 作者微調了 5 個 epoch 的所有資料集，batch size 為 32，學習率為 2 × 10-5。對於 Cramming 訓練的模型這不是最優的，其可以從 16 的 batch size 和 4 × 10−5 的學習率以及餘弦衰減中獲得微小的改進（此設置不會改善預訓練的 BERT check point）。

表 3 和表 4 描述了此設定在 GLUE 下游任務上的效能。作者比較了原始的 BERT-base check point、在達到算力上限後停止的 BERT 預訓練設置、Izsak 等人 2021 年研究中描述的設置和修改後的設置，為每塊 GPU 設置訓練一天。總體而言，效能出奇地好，尤其是對於 MNLI、QQP、QNLI 和 SST-2 等較大的資料集，下游微調可以消除完整 BERT 模型和 Cramming 設定變體之間的剩餘差異。 

此外，作者發現新方法與算力有限的普通 BERT 訓練及 Izsak 等人所描述的方法相比都有很大改進。對於Izsak 等人的研究，其描述的方法最初是為一個完整的8 GPU 刀鋒伺服器設計的，並且在新的場景中，將其中的BERT-large 模型壓縮到較小的GPU 上是導致大部分效能下降的原因。

表格3：基線BERT 與Cramming 版本模型的GLUE-dev 效能比較

其中所有任務的超參數都是固定的，epoch 限制為5 個，缺失值為NaN。是為 8 GPU 刀鋒伺服器設計的，而在這裡，所有運算都塞進了一塊 GPU。

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

整體而言，使用論文中的方法，訓練結果已經非常接近原版BERT，但要知道後者使用的總FLOPS 是新方法45-136 倍（在16 塊TPU 上要花費四天時間）。而當訓練時間延長 16 倍時（在 8 塊 GPU 上訓練兩天），新方法的性能實際上比原始 BERT 提高了很多，達到了 RoBERTa 的水平。

總結

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

以上是一日一卡挑戰：RTX2080Ti搞定大模型訓練，算力省136倍，學界歡呼的詳細內容。更多資訊請關注PHP中文網其他相關文章！