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

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

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

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

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

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

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

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

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

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

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

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

2. GPU 針對大而寬的模型進行了優化

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

速度估計

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

寬度

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

深度

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

Batch size

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

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

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

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

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

3. 語言建模任務演示: Wikitext-103

現在我們已經知道了模型尺寸和訓練速度之間的關系，我們可以依此預測: 對于給定的 GPU 時間或預算，適合目標任務的最佳模型尺寸及其能達到的性能。

這里使用的價格是 Google 云平臺 (Google Cloud Platform，GCP) 的價格。我們使用了 Peter Henderson 的 Experiment impact tracker 來估算能源消耗，并使用了 Electricity map 的荷蘭數據 (Google 的歐洲服務器所在地) 來估算 CO2 排放量。盡管巨大的訓練成本常常博得頭條，但事實上，我們仍然有可能以 30 美元的價格在中等規模的數據集上復現最先進的結果！對于一個恰當優化過的訓練方案而言，V100 已經算一個強大的武器了。

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

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

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

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

編輯：黃飛

?