獨家 | 如何在GPU資源受限情況下微調超大模型

翻譯：陳之炎

校對：趙茹萱

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

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

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

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

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

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

• 簡(jiǎn)單模式：無(wú)法適配批大小為1
• 專(zhuān)業(yè)模式：參數也沒(méi)辦法適配

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

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

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

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

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

在學(xué)習了梯度檢查點(diǎn)的細節之后，來(lái)看看如何在PyTorch中應用這個(gè)概念，看起來(lái)并不太難： 

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

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

沒(méi)有梯度累積的標準訓練環(huán)通常為： 

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

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

某些網(wǎng)絡(luò )體系結構使用專(zhuān)用的批處理操作，如BatchNorm，當使用相同的批大小時(shí)，結果可能會(huì )略有不同。

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

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

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

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

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

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

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

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

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

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

把以上方法全部利用起來(lái)，查看一下代碼： 

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

在本博中，給出了高效使用內存的關(guān)鍵概念，它適用于多種艱巨的任務(wù)，如上文所述。
將在后續的文章中討論其他概念。
衷心感謝，撥冗閱讀本文!
原文標題：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