提問：模型大小超過GPU 容量怎么辦？ 

本文的靈感來自于Yandex數(shù)據(jù)分析學院教授的“高效深度學習系統(tǒng)”課程。

預備知識：假設讀者已經(jīng)了解神經(jīng)網(wǎng)絡的前傳遞和后向傳遞的工作原理，這對理解本文內(nèi)容至關(guān)重要。文中使用PyTorch作為框架。

開始吧！

當試圖使用大型模型(即aka gpt-2-xl)，它帶有 5億多個參數(shù)，而你的GPU 資源受限，無法將它安裝到GPU上運行，或者在模型訓練期間無法實現(xiàn)論文中定義的批大小，此時該怎么辦？也許可以選擇放棄，使用一個更輕量級版本的模型，或者減小訓練的批大小，這樣的話，便無法獲得論文中描述的訓練結(jié)果。

但是，有一些技術(shù)可以幫助解決上述問題。

下面來討論一些方法，即如何利用這些方法來微調(diào)帶有15億個參數(shù)的GPT-2-XL模型。

問題的核心

首先，來了解一下將模型加載到GPU中所需GPU內(nèi)存問題的實質(zhì)。

假設模型具有 個FP32（32位浮點）參數(shù)，需要在GPU上訓練這個模型，例如，運行Adam優(yōu)化器。

通過計算，結(jié)果令人震驚。

假設已有一塊帶有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)存，由于還沒有開始訓練，也沒有加載優(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)容：

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

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

開始

提問：模型比GPU容量大，怎么辦？

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

概述

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

神經(jīng)網(wǎng)絡層的激活對應于用f標記的節(jié)點，在正向傳遞期間，按順序?qū)λ羞@些節(jié)點進行計算。對應于這些層的激活和參數(shù)的損失梯度用b標記的節(jié)點表示。在反向傳遞期間，所有這些節(jié)點都以相反的順序進行計算。f個節(jié)點的計算結(jié)果用于計算b個節(jié)點，因此所有f個節(jié)點在向前傳遞后都保存在內(nèi)存中。只有當反向傳播進展到足夠計算出f節(jié)點的所有依賴關(guān)系時，它才能從內(nèi)存中擦除。這意味著：簡單的反向傳播所需的內(nèi)存隨神經(jīng)網(wǎng)絡層數(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次順序重新計算。由于計算速度較慢，這種方法并不適用于深度學習。

為了在內(nèi)存和計算之間取得平衡，需要提出一種策略，允許重新計算節(jié)點，但次數(shù)不要太頻繁。在這里使用這樣一種策略：將神經(jīng)網(wǎng)絡激活的一個子集標記為檢查點節(jié)點。 

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

例程：

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

梯度累積/微批次

概述

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

什么是梯度累積？

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

例程：

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

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

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

重點

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

混合精度訓練

概述?

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

主要優(yōu)勢

混合精度訓練的主要優(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 來處理這些特殊情況。

8位優(yōu)化器

增加模型尺寸是獲得更佳性能的有效途徑。然而，訓練大模型時需要存儲模型、梯度和優(yōu)化器的狀態(tài)(例如，Adam的指數(shù)平滑和及先前梯度的平方和)，所有這些都存儲在數(shù)量有限的可用內(nèi)存之中。

將32位優(yōu)化器降到8位優(yōu)化器，將數(shù)值的范圍從232減少到僅2?=256，會對優(yōu)化器預留的內(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)換，無需慢速復制到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)存，看起來相當棒！

理解了它的可用性之后，再來看看如何用python實現(xiàn)它。

由Facebook提供的Bitsandbytes 包是一個圍繞CUDA自定義函數(shù)的輕量級包裝器，封裝了 8位優(yōu)化器和量化函數(shù)，利用它可以實現(xiàn)8位Adam的使用。

例程：

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

綜合上述全部方法，對GPU上的GPT-2-XL進行微調(diào)。

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

• 梯度檢查點；
• 混合精度訓練(我設了一個技巧：使用相同模型的兩個樣本。首先，用.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)，絕了！

結(jié)論

在本博中，給出了高效使用內(nèi)存的關(guān)鍵概念，它適用于多種艱巨的任務，如上文所述。將在后續(xù)的文章中討論其他概念。衷心感謝，撥冗閱讀本文!