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

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

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

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

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

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

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

體驗(yàn)地址: https://hf.co/calculator/

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

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

體驗(yàn)地址: https://hf.co/calculator/

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

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

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

速度估計(jì)

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

寬度

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

深度

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

Batch size

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

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

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

最終，我們得到運(yùn)行速度的估算公式如下:

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

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

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

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

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

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

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

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

編輯：黃飛

?