翻譯：陳之炎

校對：趙茹萱

提問：模型大小超過GPU 容量怎么辦？ 
本文的靈感來自于Yandex數(shù)據(jù)分析學(xué)院教授的“高效深度學(xué)習(xí)系統(tǒng)”課程。
預(yù)備知識：假設(shè)讀者已經(jīng)了解神經(jīng)網(wǎng)絡(luò)的前傳遞和后向傳遞的工作原理，這對理解本文內(nèi)容至關(guān)重要。文中使用PyTorch作為框架。

當(dāng)試圖使用大型模型(即aka gpt-2-xl)，它帶有 5億多個參數(shù)，而你的GPU 資源受限，無法將它安裝到GPU上運行，或者在模型訓(xùn)練期間無法實現(xiàn)論文中定義的批大小，此時該怎么辦？也許可以選擇放棄，使用一個更輕量級版本的模型，或者減小訓(xùn)練的批大小，這樣的話，便無法獲得論文中描述的訓(xùn)練結(jié)果。
但是，有一些技術(shù)可以幫助解決上述問題。
下面來討論一些方法，即如何利用這些方法來微調(diào)帶有15億個參數(shù)的GPT-2-XL模型。

首先，來了解一下將模型加載到GPU中所需GPU內(nèi)存問題的實質(zhì)。
假設(shè)模型具有 個FP32（32位浮點）參數(shù)，需要在GPU上訓(xùn)練這個模型，例如，運行Adam優(yōu)化器。
通過計算，結(jié)果令人震驚。 

假設(shè)已有一塊帶有12 GB內(nèi)存的NVIDIA GeForce RTX 3060。首先， 1e9個FP32參數(shù)約占4 GB的GPU內(nèi)存。同樣，對于梯度，也將保留相同數(shù)量的內(nèi)存。所以，總共已經(jīng)保留了8 GB的內(nèi)存，由于還沒有開始訓(xùn)練，也沒有加載優(yōu)化器，加載優(yōu)化器也同樣需要一定數(shù)量的內(nèi)存。Adam優(yōu)化器需要為每個參數(shù)存儲第一備份和第二備份，即需要8 GB額外內(nèi)存。
算下來，必須有大約16 GB的GPU內(nèi)存，才能正確地將模型加載到GPU上，在本文的例子中，GPU只有12 GB的空閑內(nèi)存?？雌饋砗懿幻?，對吧？
然而，可以通過一些方法來嘗試解決這個問題，以下是相關(guān)內(nèi)容：

• 梯度積累/微批量；
• 梯度檢查點；
• 模型并行訓(xùn)練；
• 管道作業(yè)；
• 張量并行化
• 混合精度訓(xùn)練；
• 內(nèi)存卸載；
• 優(yōu)化器8位量化。

接下來，將詳細解讀這些技術(shù)。

• 簡單模式：無法適配批大小為1
• 專業(yè)模式：參數(shù)也沒辦法適配

如果模型大于GPU容量，即便將批大小設(shè)為1都不夠，那該怎么辦呢？有一個解決方案，即設(shè)置梯度檢查點，下面來看看這個概念。
對于一個簡單的包含n層的前饋神經(jīng)網(wǎng)絡(luò)來說，梯度的計算圖如下： 

神經(jīng)網(wǎng)絡(luò)層的激活對應(yīng)于用f標(biāo)記的節(jié)點，在正向傳遞期間，按順序?qū)λ羞@些節(jié)點進行計算。對應(yīng)于這些層的激活和參數(shù)的損失梯度用b標(biāo)記的節(jié)點表示。在反向傳遞期間，所有這些節(jié)點都以相反的順序進行計算。f個節(jié)點的計算結(jié)果用于計算b個節(jié)點，因此所有f個節(jié)點在向前傳遞后都保存在內(nèi)存中。只有當(dāng)反向傳播進展到足夠計算出f節(jié)點的所有依賴關(guān)系時，它才能從內(nèi)存中擦除。這意味著：簡單的反向傳播所需的內(nèi)存隨神經(jīng)網(wǎng)絡(luò)層數(shù)n的變化呈線性增長。
下面是這些節(jié)點的計算順序，紫色陰影圓圈表示在給定時間里需要將哪個節(jié)點保存到內(nèi)存之中。 

如上所述的簡單反向傳播在計算方面是最優(yōu)的：它只計算每個節(jié)點一次。但是，如果重新計算節(jié)點，可能會節(jié)省大量內(nèi)存。例如，可以簡單地重新計算每個節(jié)點。執(zhí)行的順序和所使用的內(nèi)存如下圖所示： 

這種策略在內(nèi)存方面是最優(yōu)的。但是，請注意，節(jié)點計算的數(shù)量進行了n2次縮放，而先前的縮放系數(shù)為n：每個n個節(jié)點都按n次順序重新計算。由于計算速度較慢，這種方法并不適用于深度學(xué)習(xí)。
為了在內(nèi)存和計算之間取得平衡，需要提出一種策略，允許重新計算節(jié)點，但次數(shù)不要太頻繁。在這里使用這樣一種策略：將神經(jīng)網(wǎng)絡(luò)激活的一個子集標(biāo)記為檢查點節(jié)點。 

在本示例中，選擇將第sqrt(n)個節(jié)點標(biāo)記為檢查點。這樣，檢查點節(jié)點的數(shù)量和檢查點之間的節(jié)點數(shù)量都在sqrt(n)之間，這意味著：所需的內(nèi)存量也按n的順序進行了縮放。該策略所需的額外計算量相當(dāng)于網(wǎng)絡(luò)單次前向傳遞所需的計算量。

在學(xué)習(xí)了梯度檢查點的細節(jié)之后，來看看如何在PyTorch中應(yīng)用這個概念，看起來并不太難： 

深度學(xué)習(xí)模型正在越變越大，很難在GPU內(nèi)存中安裝這樣大型的神經(jīng)網(wǎng)絡(luò)。因此，被迫在訓(xùn)練時選用較小的批大小，它可能導(dǎo)致較慢的收斂和較低的準(zhǔn)確性。

在訓(xùn)練神經(jīng)網(wǎng)絡(luò)時，通常會將數(shù)據(jù)分批量處理，神經(jīng)網(wǎng)絡(luò)預(yù)測批處理標(biāo)簽，用于計算相對于實際目標(biāo)的損失。接下來，執(zhí)行反向傳遞計算出梯度，更新模型權(quán)值。
梯度累積對訓(xùn)練過程的最后一步進行了修正：在繼續(xù)下一個小批之前，保存梯度值，并將新的梯度添加到之前保存的梯度中，用這種方法取代更新每個小批的網(wǎng)絡(luò)權(quán)重。只有在模型處理了幾個小批次后，才會更新權(quán)重。
梯度積累模擬了一個更大的批大小，如果想在一個小批中使用64張圖像，如果批大小超過了8，則會報“CUDA內(nèi)存出錯…”。在這種情況下，可以使用8批圖像，并在模型處理64/8=8批后更新一次權(quán)重。如果你從這8個批次中積累每一個梯度，結(jié)果將是（幾乎）相同的，這樣便能夠執(zhí)行訓(xùn)練啦！

沒有梯度累積的標(biāo)準(zhǔn)訓(xùn)練環(huán)通常為： 

在PyTorch中，梯度累積可以很容易地完成。模型利用accumulation_steps處理完成小批之后，便可以執(zhí)行優(yōu)化。還可以利用accumulation_steps根據(jù)損失函數(shù)的性質(zhì)來劃分運行損失： 

真漂亮，對嗎？當(dāng)調(diào)用loss.backward() 時計算梯度，并由PyTorch累積，直到調(diào)用optimizer.zero_grad()時停止。

某些網(wǎng)絡(luò)體系結(jié)構(gòu)使用專用的批處理操作，如BatchNorm，當(dāng)使用相同的批大小時，結(jié)果可能會略有不同。

混合精度訓(xùn)練是指將部分或全部FP32參數(shù)轉(zhuǎn)換為更小的格式，如FP16、TF16（浮點張量）或BF16(浮點字節(jié))。

混合精度訓(xùn)練的主要優(yōu)勢是：

• 減少內(nèi)存使用；
• 性能提速（更高的算術(shù)強度或更小的通信占用）；
• 使用專用硬件進行更快地計算。

目前只對第一個優(yōu)勢感興趣——減少內(nèi)存的使用量，來看看如何使用PyTorch模型實現(xiàn)它。

結(jié)果，在完成.half()操作之后，模型變小了2倍。
將模型轉(zhuǎn)換為不同的格式(即BF16，TF16)后的縮放損失，將在后續(xù)的文章中討論。
有些操作在FP16中是無法完成的，如Softmax。PyTorch可利用torch.autocast 來處理這些特殊情況。

增加模型尺寸是獲得更佳性能的有效途徑。然而，訓(xùn)練大模型時需要存儲模型、梯度和優(yōu)化器的狀態(tài)(例如，Adam的指數(shù)平滑和及先前梯度的平方和)，所有這些都存儲在數(shù)量有限的可用內(nèi)存之中。
將32位優(yōu)化器降到8位優(yōu)化器，將數(shù)值的范圍從232減少到僅2?=256，會對優(yōu)化器預(yù)留的內(nèi)存數(shù)量產(chǎn)生巨大的影響。
研究人員提出了一種新的8位Adam優(yōu)化器，論文作者在文中這么說： “它將32位的性能維持到部分原始內(nèi)存中”。
8位優(yōu)化器有三個組成部分：（1）塊級量化，隔離異常值，將誤差均勻地分配給每一個比特；（2）動態(tài)量化，高精度地量化小值和大值；（3）穩(wěn)定的嵌入層，以提高詞嵌入優(yōu)化模型的穩(wěn)定性。
有了這些組件，可直接使用8位狀態(tài)執(zhí)行優(yōu)化。將8位優(yōu)化器狀態(tài)量化為32位，執(zhí)行更新，然后再將狀態(tài)量化為8位進行存儲。在寄存器中逐元素進行8位到32位的轉(zhuǎn)換，無需慢速復(fù)制到GPU內(nèi)存或額外的臨時內(nèi)存中執(zhí)行量化和去量化。對于GPU來說，這意味著8位優(yōu)化器要快于常規(guī)的32位優(yōu)化器。
來看看使用8位Adam之后，鼓舞人心的結(jié)果： 

可以看出，使用量化的Adam可以節(jié)省大約8.5 GB的GPU內(nèi)存，看起來相當(dāng)棒！
理解了它的可用性之后，再來看看如何用python實現(xiàn)它。
由Facebook提供的Bitsandbytes 包是一個圍繞CUDA自定義函數(shù)的輕量級包裝器，封裝了 8位優(yōu)化器和量化函數(shù)，利用它可以實現(xiàn)8位Adam的使用。

如上所述，量化優(yōu)化器的使用非常簡單，結(jié)果也不錯。

最后，在掌握了上述方法之后，利用這些方法來解決實際問題，對擁有15億個參數(shù)的GPT-2-XL模型進行微調(diào)。顯然，無法將它加載到12 GB內(nèi)存的NVIDIA GeForce RTX 3060 GPU之上。
列出可以使用的全部方法：

• 梯度檢查點；
• 混合精度訓(xùn)練(我設(shè)了一個技巧：使用相同模型的兩個樣本。首先，用.half將它加載到GPU上，將其命名為gpu_model；其次，在CPU上，將其命名為cpu_model。評估好GPU模型之后，將 gpu_model的梯度加載到cpu_model中，運行optimizer.step()，將更新后的參數(shù)加載到gpu_model上)；
• 使用batch_size=64，minibatch_size=4的梯度累積，需要通過 accumulation_steps來縮放損失；
• 8位Adam優(yōu)化器。

把以上方法全部利用起來，查看一下代碼： 

利用上述所有方法之后，在GPU上實現(xiàn)了對16GB的GPT-2-XL模型微調(diào)，絕了！

在本博中，給出了高效使用內(nèi)存的關(guān)鍵概念，它適用于多種艱巨的任務(wù)，如上文所述。
將在后續(xù)的文章中討論其他概念。
衷心感謝，撥冗閱讀本文!
原文標(biāo)題：How to fine tune VERY large model if it doesn’t fit on your GPU原文鏈接：https://medium.com/@bestasoff/how-to-fine-tune-very-large-model-if-it-doesnt-fit-on-your-gpu-3561e50859af文：Stanislav Belyasov