在K40小破卡訓(xùn)練50層BERT Large的寶藏trick

夕小瑤?機(jī)器學(xué)習(xí) 話題的優(yōu)秀回答者434 人贊同了該文章

前言

雖然TPU的顯存令人羨慕，但是由于眾所周知的原因，絕大部分人還是很難日?；褂玫?。英偉達(dá)又一直在擠牙膏，至今單卡的最大顯存也僅僅到32G（參考V100、DGX-2）。然而，訓(xùn)練一個(gè)24層的BERT Large模型的時(shí)候，如果sequence length開(kāi)滿512，那么batch size僅僅開(kāi)到8（有時(shí)候能到10）就把這寥寥32G的顯存打滿了。如果想訓(xùn)練一個(gè)48層乃至100層的BERT Large，那完全是土豪們的游戲了，需要瘋狂的模型并行+分布式多機(jī)訓(xùn)練。

但！是！萬(wàn)能的小夕前不久在Daxiang Dong大佬的安利下，發(fā)現(xiàn)了

@陳天奇 大佬2016年的一篇寶藏paper！

傳送門(mén)：https://arxiv.org/pdf/1604.06174.pdf

簡(jiǎn)單的劃一下重點(diǎn)：

這篇paper用時(shí)間換空間的思想，在前向時(shí)只保存部分中間節(jié)點(diǎn)，在反向時(shí)重新計(jì)算沒(méi)保存的部分。論文通過(guò)這種機(jī)制，在每個(gè)batch只多計(jì)算一次前向的情況下，把n層網(wǎng)絡(luò)的占用顯存優(yōu)化到了 ( √)O(n)。在極端情況下，仍可用 ( )O(nlogn)的計(jì)算時(shí)間換取到 ( )O(logn)的顯存占用。在論文的實(shí)驗(yàn)中，他們成功將將1000層的殘差網(wǎng)絡(luò)從48G優(yōu)化到了7G。且，這種方法同樣可以直接應(yīng)用于RNN結(jié)構(gòu)中。

看完摘要，瞬間感覺(jué)在小破卡上訓(xùn)練BERT Large有救了！！！

此外，來(lái)快速過(guò)一遍paper中最重要的三點(diǎn)結(jié)論：

1. 梯度計(jì)算等價(jià)，理論上沒(méi)有精度損失

2. 可以節(jié)省4倍+的顯存開(kāi)銷(xiāo)

3. 訓(xùn)練速度僅僅會(huì)被拖慢30%

不過(guò)論文發(fā)表在2016年，當(dāng)時(shí)還沒(méi)有BERT，不過(guò)Baidu Paddle團(tuán)隊(duì)補(bǔ)了一個(gè)BERT的實(shí)驗(yàn)結(jié)果，發(fā)現(xiàn)在BERT上面只用22.5%的訓(xùn)練速度損失就能換來(lái)5倍+的顯存開(kāi)銷(xiāo)節(jié)省！相關(guān)實(shí)驗(yàn)在本文末尾，不著急，接下來(lái)我們先一起分析一下在訓(xùn)練階段時(shí)顯存為什么容易不足。

感謝Baidu Paddle團(tuán)隊(duì)提供本節(jié)圖文素材和測(cè)試數(shù)據(jù)

訓(xùn)練階段顯存為何不足

深度學(xué)習(xí)中，網(wǎng)絡(luò)的一次訓(xùn)練包含前向計(jì)算、后向計(jì)算和優(yōu)化三個(gè)步驟。

在這個(gè)過(guò)程中，前向計(jì)算會(huì)輸出大量的隱層變量Tensor，當(dāng)模型層數(shù)加深時(shí)，Tensor數(shù)量可達(dá)成千上萬(wàn)個(gè)。如Bert Large模型，單個(gè)Tensor可達(dá)到1GB，這些Tensor在顯存中累積，顯存很快就爆掉了

下圖是Bert Large模型在一次訓(xùn)練過(guò)程中的顯存使用情況，可以明顯看到在前向計(jì)算過(guò)程中，顯存累積趨勢(shì)是一個(gè)陡峭的上升直線。而在反向計(jì)算過(guò)程中，這些隱層Tensor又會(huì)很快地被消耗掉，又是一個(gè)陡峭的下降曲線，顯存直接降到低位。

那么問(wèn)題來(lái)了，為什么不直接刪除這些前向計(jì)算的Tensor呢？

答案很簡(jiǎn)單，因?yàn)檫@些隱層的Tensor在反向的時(shí)會(huì)被用到（手動(dòng)狗頭

來(lái)個(gè)簡(jiǎn)單的證明。

假設(shè)前向計(jì)算中有一個(gè)矩陣乘法計(jì)算：

Y = W × X

對(duì)W求梯度：

很容易發(fā)現(xiàn)，對(duì)W求梯度的公式里有X，而X就是那個(gè)巨能吃顯存的隱層Tensor！

那我們是否可以暫時(shí)扔掉這些隱層Tensor，在反向計(jì)算時(shí)再把它們重新生成出來(lái)呢？當(dāng)然可以，這正是上面這篇paper的思想。

重計(jì)算

顧名思義，"重計(jì)算"就是讓每個(gè)訓(xùn)練迭代過(guò)程做兩次前向計(jì)算，看起來(lái)有點(diǎn)奇怪，實(shí)際上卻非常有效！對(duì)于剛剛那個(gè)吃顯存的Bert Large，支持重計(jì)算機(jī)制后，顯存占用直接從175GB降低到20GB，陡峭的顯存上升直線變成了緩慢增長(zhǎng)的Z形曲線，如下圖所示。

核心思想是將前向計(jì)算分割成多個(gè)段，將每個(gè)段的起始Tensor作為這個(gè)段的檢查點(diǎn)（checkpoints）。前向計(jì)算時(shí)，除了檢查點(diǎn)以外的其他隱層Tensor占有的顯存可以及時(shí)釋放。反向計(jì)算用到這些隱層Tensor時(shí)，從前一個(gè)檢查點(diǎn)開(kāi)始，重新進(jìn)行這個(gè)段的前向計(jì)算，就可以重新獲得隱層Tensor。

重計(jì)算機(jī)制有點(diǎn)像玩單機(jī)游戲。每過(guò)一個(gè)關(guān)卡就會(huì)保存一個(gè)檢查點(diǎn)，而隱層Tensor就相當(dāng)于游戲中任何一個(gè)時(shí)刻的圖像。普通的訓(xùn)練方式是打通一遍游戲，并且將游戲中所有時(shí)刻的圖像保存下來(lái)；而重計(jì)算機(jī)制的思路是先把游戲通關(guān)，保存檢查點(diǎn)，后面當(dāng)收到某一時(shí)刻圖像的請(qǐng)求時(shí)，再重打一遍這一關(guān)卡就可以了。

如下圖，舉一個(gè)簡(jiǎn)單的例子，添加重計(jì)算機(jī)制前，前向計(jì)算中需要存儲(chǔ)的隱層是4個(gè)紅點(diǎn)；添加重計(jì)算機(jī)制后，需要存儲(chǔ)的隱層變?yōu)?個(gè)藍(lán)點(diǎn)， 從而節(jié)省了這部分內(nèi)存。

雖然時(shí)間也是寶貴的，但重計(jì)算方法的性價(jià)比很高。在論文的實(shí)驗(yàn)中，作者用30%的計(jì)算時(shí)間換取了4倍的內(nèi)存空間。并且重計(jì)算只是重復(fù)了一次前向的過(guò)程，理論上精度沒(méi)有任何損失。

這么寶藏的算法當(dāng)然也少不了開(kāi)源實(shí)現(xiàn)。

開(kāi)源實(shí)現(xiàn)

調(diào)研了一波，似乎TF沒(méi)有原生支持，但是生態(tài)里有OpenAI的第三方實(shí)現(xiàn)；pytorch和paddlepaddle中都有原生API支持

• Pytorch：
• • torch.utils.checkpoint
• PaddlePaddle：
• • optimizer.RecomputeOptimizer

不過(guò)pytorch的文檔比較略，也沒(méi)有提供更細(xì)致的示例和相關(guān)數(shù)據(jù)，有興趣的小伙伴自行試一下。paddle框架中提供了詳細(xì)到哭的文檔，甚至還有一個(gè)現(xiàn)成的BERT+重計(jì)算的例子，以及非常詳細(xì)的實(shí)驗(yàn)測(cè)試結(jié)果。這里直接貼過(guò)來(lái)（真香系列

Paddle中實(shí)現(xiàn)顯存重計(jì)算大體分為三步：

1. 定義一個(gè)經(jīng)典的優(yōu)化器，如SGD優(yōu)化器；
2. 在外面包一層重計(jì)算優(yōu)化器；
3. 設(shè)置檢查點(diǎn)。

以MLP為例，只需要增加兩行代碼就可以進(jìn)入重計(jì)算模式

import paddle.fluid as fluid
# 定義MLP
def mlp(input_x, input_y, hid_dim=128, label_dim=2):print(input_x)fc_1 = fluid.layers.fc(input=input_x, size=hid_dim)prediction = fluid.layers.fc(input=[fc_1], size=label_dim, act='softmax')cost = fluid.layers.cross_entropy(input=prediction, label=input_y)sum_cost = fluid.layers.reduce_mean(cost)return sum_cost, fc_1, prediction

總結(jié)

以上是生活随笔為你收集整理的小显卡训练大模型的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。