在這個全民 LLM 的狂歡里，想測測你拿到的預(yù)算夠訓(xùn)一個多大的模型嗎？本文會給你一個答案，至少給你一個計算公式。

在自然語言處理領(lǐng)域，有時候我們恍惚覺得大家是為了搏頭條而在模型尺寸上不斷進(jìn)行軍備競賽。1750 億參數(shù) 無疑是一個很抓眼球數(shù)字！為什么不考慮高效地去訓(xùn)一個小一點的模型？其實，這是因為深度學(xué)習(xí)領(lǐng)域有一個挺驚人的縮放效應(yīng)，那就是: 大神經(jīng)網(wǎng)絡(luò)計算效率更高。這是以 OpenAI 為代表的團(tuán)隊在像 神經(jīng)語言模型的縮放定律 這樣的論文中探索出的結(jié)論。本文的研究也基于這一現(xiàn)象，我們將其與 GPU 速度估計相結(jié)合，用于確保在進(jìn)行語言模型實驗時，我們能根據(jù)我們算力預(yù)算來設(shè)計最合適的模型尺寸 (劇透一下，這個大小比你想象的要大！)。我們將展示我們的方法是如何影響一個標(biāo)準(zhǔn)的語言建?；鶞?zhǔn)的架構(gòu)決策的: 我們在沒有任何超參優(yōu)化的前提下，僅使用了原論文 75% 的訓(xùn)練時間，復(fù)現(xiàn)了 Zhang 等人的 Transformer-XL 論文 中的 14 層模型的最佳結(jié)果。我們還估計 來自同一篇論文的 18 層模型其實僅需要比原論文少一個數(shù)量級的步數(shù)就能達(dá)到相同的結(jié)果。繼續(xù)閱讀之前想先玩玩我們的演示嗎？只需閱讀原文，查看交互式示例！

“1750 億參數(shù)”出處論文:
https://arxiv.org/abs/2005.14165

神經(jīng)語言模型的縮放定律論文:
https://arxiv.org/abs/2001.08361

Transformer-XL 論文:
https://arxiv.org/pdf/1901.02860.pdf

1. 停止訓(xùn)練的最佳時間 (比你想象的要早)

我們先觀察一些損失曲線 (loss curve)。我們使用的任務(wù)是在 Wikitext-103 上訓(xùn)練 Transformer-XL 語言模型，Wikitext-103 是一個標(biāo)準(zhǔn)的中等體量的測試基準(zhǔn)。GPT-2 在此等體量的數(shù)據(jù)集上表現(xiàn)不佳。隨著訓(xùn)練的進(jìn)行，我們來觀察計算成本 (通過浮點運算數(shù)來衡量) 與模型性能 (通過驗證集上的損失來衡量) 的聯(lián)動關(guān)系。我們做點實驗吧！在下圖中，不同顏色的線段表示不同層數(shù)和大小的 Transformer-XL 模型運行 200000 步的數(shù)據(jù)，這些模型除了層數(shù)與大小外的所有其他超參數(shù)都相同。模型參數(shù)量范圍從幾千到一億 (不含嵌入)。越大的模型在圖中越靠右，因為它們每一步需要的計算量更大。本圖是交互式的，你可以玩一玩！

體驗地址: https://hf.co/calculator/

與 縮放定律 一文中的做法一樣，我們的橫軸為非嵌入浮點運算數(shù) (non-embedding FLoating Point Operations, neFLOs)，縱軸為驗證集損失。對于給定的 neFLOs 預(yù)算，似乎存在一個任何模型都沒法越過的性能邊界，我們在圖中用紅色線段表示。在 縮放定律 一文中，它被稱為計算邊界 (compute frontier)。我們可以看到，在所有的實驗中，幾乎每個實驗都能在經(jīng)過初始若干步的損失迅速降低后達(dá)到或接近該計算邊界，隨后又在訓(xùn)練接近尾聲時，因訓(xùn)練效率降低而偏離該計算邊界。這個現(xiàn)象有其實際意義: 給定固定的浮點運算預(yù)算，為了達(dá)到最佳性能，你應(yīng)該選擇一個模型尺寸，使得在浮點運算預(yù)算見頂時正好達(dá)到計算邊界，然后我們就可以在此時停止訓(xùn)練。此時離模型收斂所需的時間還很遠(yuǎn)，模型收斂還需要 10 倍左右的時間。事實上，如果此時你還有額外的錢可用于計算，你應(yīng)該把大部分用到增大模型上，而只將一小部分用于增加訓(xùn)練步數(shù)。[ 譯者注: 這是因為性能邊界本質(zhì)上度量了每 neFLOs 帶來的 loss 的降低是多少，到達(dá)計算邊界后，后面的每 neFLOs 能帶來的 loss 的降低變小，不劃算了。我們應(yīng)該轉(zhuǎn)而去尋求增大模型所帶來的接近計算邊界的高回報，而不應(yīng)該卷在增加訓(xùn)練步數(shù)帶來的低回報上。 ]

在 縮放定律 一文中，OpenAI 團(tuán)隊用冪律函數(shù)擬合了一個 GPT-2 訓(xùn)練的計算邊界。這似乎也適用于我們的任務(wù)，我們也擬合了一個以預(yù)算為自變量，最適合該預(yù)算的模型參數(shù)量為因變量的冪律函數(shù)。如下圖所示。

體驗地址: https://hf.co/calculator/

由于好的模型的 neFLOs- 損失曲線往往會與計算邊界相切比較長時間，因此最終的擬合函數(shù)會存在一些噪聲。然而，這恰恰也意味著基于該擬合函數(shù)的估計的容忍度會比較好，即使我們預(yù)測尺寸有點偏差，其結(jié)果仍然會非常接近最優(yōu)值。我們發(fā)現(xiàn)，如果將算力預(yù)算乘以 10，最佳模型尺寸會隨之乘以 7.41，而最佳訓(xùn)練步數(shù)僅乘以 1.35。將此規(guī)則外推到 Tranformer-XL 論文中的那個更大的 18 層最先進(jìn)模型，我們發(fā)現(xiàn) 其最佳訓(xùn)練步數(shù)約為 25 萬步。即使這個數(shù)字由于模型尺寸的變化而變得不那么精確，它也比論文中所述的 收斂所需的 400 萬步 小得多。以更大的模型、更少的訓(xùn)練步數(shù)為起點，在給定的 (龐大的) 預(yù)算下我們能訓(xùn)到更小的損失。

2. GPU 針對大而寬的模型進(jìn)行了優(yōu)化

我們現(xiàn)在有了一個將性能和最佳模型尺寸與 neFLOs 聯(lián)系起來的規(guī)則。然而，neFLOs 有點難以具象化。我們能否將其轉(zhuǎn)化為更直觀的指標(biāo)，如訓(xùn)練時間？其實，無論你是有時間上的限制還是財務(wù)上的限制，主要關(guān)注的都是 GPU 時間。為了在 neFLOs 和 GPU 時間之間建立聯(lián)系，我們在谷歌云平臺上用 4 種不同 GPU 實例以及各種不同大小的 Transformer-XL 模型進(jìn)行了數(shù)萬次的基準(zhǔn)測試 (包括混合精度訓(xùn)練測試)。以下是我們的發(fā)現(xiàn):

速度估計

每秒 neFLOs (即公式中的 speed) 可以建模為由模型寬度 (每層神經(jīng)元數(shù)) 、深度 (層數(shù)) 和 batch size 三個因子組成的多變量函數(shù)，這三個因子的重要性遞減。在我們的實驗中，觀察到的最大預(yù)測誤差為測量值的 15%。

寬度

GPU 針對寬 transfomer 模型的大型前饋層進(jìn)行了優(yōu)化。在我們所有的實驗中，每秒 neFLOs 與模型寬度 呈 1.6 次方的冪律關(guān)系，這意味著兩倍寬的模型需要 4 倍的操作。然而執(zhí)行這些操作的速度也提高了大約 3.16 倍， 幾乎抵消了額外的計算成本。

深度

每秒 neFLOs 也與深度正相關(guān)。我們目前發(fā)現(xiàn)的最佳關(guān)系是每秒 neFLOs 與 ? 成正比。這與 transformer 模型必須串行地處理每一層的事實是一致的。從本質(zhì)上講， 層數(shù)更多的模型其實并不會更快，但它們似乎表現(xiàn)出更快，其原因主要是它們的均攤開銷更小。公式中的 常數(shù) 就代表這一開銷，在我們的實驗中該常數(shù)一直在 5 左右，這其實意味著 GPU 加載數(shù)據(jù)、嵌入和 softmax 這些操作的耗時大約相當(dāng)于 5 個 transfomer 層的時間。

Batch size

Batch size 發(fā)揮的作用最小。Batch size 較小時，其與速度呈正相關(guān)關(guān)系，但這個關(guān)系很快就飽和了(甚至在 V100 和 P100 上 batch size 大于 64 后、在 K80 和 P4 batch size 大于 16 后，速度比小 batch size 時還有所降低)。因此，我們將其對速度的貢獻(xiàn)建模為對數(shù)函數(shù)以簡化計算，它是 3 個因子中最弱的。因此，最終我們所有實驗都是在單 GPU 上用 batch size 64 運行出來的。這是大模型的另一個好處: 因為更大的 batch size 似乎沒有多大幫助，如果你的模型太大而無法塞進(jìn) GPU，你可以只使用較小的 batch size 以及梯度累積技術(shù)。

2 的冪在 2020 年仍然很重要！

最后，一個令人驚訝的收獲是 寬度或 batch size 設(shè)置為 2 的冪的話其最終性能會比設(shè)為其他值高。有或沒有 Tensor Core 的 GPU 都是如此。在像 V100 這樣的 Tensor Core GPU 上，NVIDIA 建議張量形狀設(shè)置為 8 的倍數(shù); 然而，我們試驗過將其不斷加倍至 512 的倍數(shù)，性能還會繼續(xù)提高。但是，在最終擬合時我們還是只選擇擬合 2 的冪的數(shù)據(jù)，因為擬合所有數(shù)據(jù)點意味著擬合質(zhì)量會變差，而且最終的擬合結(jié)果會對采用 2 的冪情況下的速度估計得過于樂觀。但這不妨礙你去選擇最快的形狀參數(shù)。

最終，我們得到運行速度的估算公式如下:

例如，在未使用混合精度的 V100 GPU 上，k=2.21 × 10^7、a=1.66、b=5.92、c=1.33。不同的 GPU 具有不同的乘性系數(shù)，但結(jié)果很接近。

3. 語言建模任務(wù)演示: Wikitext-103

現(xiàn)在我們已經(jīng)知道了模型尺寸和訓(xùn)練速度之間的關(guān)系，我們可以依此預(yù)測: 對于給定的 GPU 時間或預(yù)算，適合目標(biāo)任務(wù)的最佳模型尺寸及其能達(dá)到的性能。

這里使用的價格是 Google 云平臺 (Google Cloud Platform，GCP) 的價格。我們使用了 Peter Henderson 的 Experiment impact tracker 來估算能源消耗，并使用了 Electricity map 的荷蘭數(shù)據(jù) (Google 的歐洲服務(wù)器所在地) 來估算 CO2 排放量。盡管巨大的訓(xùn)練成本常常博得頭條，但事實上，我們?nèi)匀挥锌赡芤?30 美元的價格在中等規(guī)模的數(shù)據(jù)集上復(fù)現(xiàn)最先進(jìn)的結(jié)果！對于一個恰當(dāng)優(yōu)化過的訓(xùn)練方案而言，V100 已經(jīng)算一個強(qiáng)大的武器了。

Experiment impact tracker:
https://github.com/Breakend/experiment-impact-tracker

Electricity map:
https://www.electricitymap.org

圖中所示的數(shù)據(jù)的測例為在 Wikitext-103 上使用 batch size 60 以及單 GPU 訓(xùn)練一個 Transformer-XL 模型，模型的目標(biāo)長度 (target length) 和記憶長度 (memory length) 為 150，測試基于 CMU 的 Transformer-XL 代碼庫。為了充分利用 V100 的 Tensor Core 功能，我們在該 GPU 上把 batch size 設(shè)為 64，序列長度設(shè)為 152。在我們的 模型尺寸 - 速度預(yù)測公式中，我們假設(shè)內(nèi)部前饋層維度與嵌入和注意力維度相同，并且寬深比是恒定的。Reformer 表明，這種設(shè)置有利于節(jié)省內(nèi)存。雖然 縮放定律 一文表明: 形狀不會顯著影響 GPT-2 的性能。然而，對于大模型而言，我們還是發(fā)現(xiàn)具有更大前饋層的更淺的模型的性能會更好，因此我們在圖中給出了兩種候選的模型形狀: 一個寬而淺，一個窄而深。

為了復(fù)現(xiàn)中型 Transformer-XL 預(yù)訓(xùn)練模型 (損失為 3.15) 的結(jié)果，我們調(diào)整了原模型的大小以增加的前饋維度并使之為 2 的高次冪，同時保持相同參數(shù)量。我們最終得到了一個 14 層的模型，隱藏層維度為 768 且前饋層維度為 1024。相比之下，原文中的模型是通過激進(jìn)的超參數(shù)搜索搜得的 16 層模型，形狀也很奇怪，隱藏層維度為 410 且前饋層維度為 2100。我們的實驗表明，由于我們的形狀是 2 的高次方，并且是一個更淺、更寬的模型，因此它在 NVIDIA RTX Titan 上每 batch 的速度比原模型提高了 20%。對于該模型，CMU 團(tuán)隊提供的腳本已經(jīng)非常接近最佳停止時間。最終，我們獲得了相同的性能，同時減少了 25% 的訓(xùn)練時間。最重要的是，原模型使用超參數(shù)搜索得到了對它而言更優(yōu)形狀，而我們什么也沒調(diào)，甚至連隨機(jī)種子也是直接復(fù)用的他們手調(diào)的隨機(jī)種子。由于我們使用了較小規(guī)模的訓(xùn)練來擬合縮放定律，并依此縮放定律計算所需的模型超參，因此節(jié)省參數(shù)搜索實際上可能是我們獲得的另一個也是更大的一個收益。

編輯：黃飛

?