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)練的進行，我們來觀察計算成本 (通過浮點運算數(shù)來衡量) 與模型性能 (通過驗證集上的損失來衡量) 的聯(lián)動關(guān)系。我們做點實驗吧！在下圖中，不同顏色的線段表示不同層數(shù)和大小的 Transformer-XL 模型運行 200000 步的數(shù)據(jù)，這些模型除了層數(shù)與大小外的所有其他超參數(shù)都相同。模型參數(shù)量范圍從幾千到一億 (不含嵌入)。越大的模型在圖中越靠右，因為它們每一步需要的計算量更大。本圖是交互式的，你可以玩一玩！

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

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

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

2. GPU 針對大而寬的模型進行了優(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 模型進行了數(shù)萬次的基準(zhǔn)測試 (包括混合精度訓(xùn)練測試)。以下是我們的發(fā)現(xiàn):

速度估計

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

寬度

GPU 針對寬 transfomer 模型的大型前饋層進行了優(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 時還有所降低)。因此，我們將其對速度的貢獻建模為對數(shù)函數(shù)以簡化計算，它是 3 個因子中最弱的。因此，最終我們所有實驗都是在單 GPU 上用 batch size 64 運行出來的。這是大模型的另一個好處:?因為更大的 batch size 似乎沒有多大幫助，如果你的模型太大而無法塞進 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é)果很接近。