如何在GPU資源受限情況下微調超大模型
提問：模型大小超過GPU 容量怎么辦？
本文的靈感來自于Yandex數據分析學院教授的“高效深度學習系統”課程。
預備知識：假設讀者已經了解神經網絡的前傳遞和后向傳遞的工作原理，這對理解本文內容至關重要。文中使用PyTorch作為框架。
開始吧！
當試圖使用大型模型(即aka gpt-2-xl)，它帶有 5億多個參數，而你的GPU 資源受限，無法將它安裝到GPU上運行，或者在模型訓練期間無法實現論文中定義的批大小，此時該怎么辦？也許可以選擇放棄，使用一個更輕量級版本的模型，或者減小訓練的批大小，這樣的話，便無法獲得論文中描述的訓練結果。
但是，有一些技術可以幫助解決上述問題。
下面來討論一些方法，即如何利用這些方法來微調帶有15億個參數的GPT-2-XL模型。
問題的核心
首先，來了解一下將模型加載到GPU中所需GPU內存問題的實質。
假設模型具有 個FP32（32位浮點）參數，需要在GPU上訓練這個模型，例如，運行Adam優化器。
通過計算，結果令人震驚。
假設已有一塊帶有12 GB內存的NVIDIA GeForce RTX 3060。首先， 1e9個FP32參數約占4 GB的GPU內存。
同樣，對于梯度，也將保留相同數量的內存。所以，總共已經保留了8 GB的內存，由于還沒有開始訓練，也沒有加載優化器，加載優化器也同樣需要一定數量的內存。Adam優化器需要為每個參數存儲第一備份和第二備份，即需要8 GB額外內存。算下來，必須有大約16 GB的GPU內存，才能正確地將模型加載到GPU上，在本文的例子中，GPU只有12 GB的空閑內存。看起來很不妙，對吧？
然而，可以通過一些方法來嘗試解決這個問題，以下是相關內容：
- 梯度積累/微批量；
- 梯度檢查點；
- 模型并行訓練；
- 管道作業；
- 張量并行化
- 混合精度訓練；
- 內存卸載；
- 優化器8位量化。
接下來，將詳細解讀這些技術。
開始
提問：模型比GPU容量大，怎么辦？
- 簡單模式：無法適配批大小為1
- 專業模式：參數也沒辦法適配
概述
如果模型大于GPU容量，即便將批大小設為1都不夠，那該怎么辦呢？有一個解決方案，即設置梯度檢查點，下面來看看這個概念。對于一個簡單的包含n層的前饋神經網絡來說，梯度的計算圖如下：
神經網絡層的激活對應于用f標記的節點，在正向傳遞期間，按順序對所有這些節點進行計算。對應于這些層的激活和參數的損失梯度用b標記的節點表示。在反向傳遞期間，所有這些節點都以相反的順序進行計算。f個節點的計算結果用于計算b個節點，因此所有f個節點在向前傳遞后都保存在內存中。只有當反向傳播進展到足夠計算出f節點的所有依賴關系時，它才能從內存中擦除。這意味著：簡單的反向傳播所需的內存隨神經網絡層數n的變化呈線性增長。
下面是這些節點的計算順序，紫色陰影圓圈表示在給定時間里需要將哪個節點保存到內存之中。
梯度檢查點?
如上所述的簡單反向傳播在計算方面是最優的：它只計算每個節點一次。但是，如果重新計算節點，可能會節省大量內存。例如，可以簡單地重新計算每個節點。執行的順序和所使用的內存如下圖所示：
這種策略在內存方面是最優的。但是，請注意，節點計算的數量進行了n2次縮放，而先前的縮放系數為n：每個n個節點都按n次順序重新計算。由于計算速度較慢，這種方法并不適用于深度學習。
為了在內存和計算之間取得平衡，需要提出一種策略，允許重新計算節點，但次數不要太頻繁。在這里使用這樣一種策略：將神經網絡激活的一個子集標記為檢查點節點。
在本示例中，選擇將第sqrt(n)個節點標記為檢查點。這樣，檢查點節點的數量和檢查點之間的節點數量都在sqrt(n)之間，這意味著：所需的內存量也按n的順序進行了縮放。該策略所需的額外計算量相當于網絡單次前向傳遞所需的計算量。
例程：
在學習了梯度檢查點的細節之后，來看看如何在PyTorch中應用這個概念，看起來并不太難：
梯度累積/微批次
概述
深度學習模型正在越變越大，很難在GPU內存中安裝這樣大型的神經網絡。因此，被迫在訓練時選用較小的批大小，它可能導致較慢的收斂和較低的準確性。
什么是梯度累積？
在訓練神經網絡時，通常會將數據分批量處理，神經網絡預測批處理標簽，用于計算相對于實際目標的損失。接下來，執行反向傳遞計算出梯度，更新模型權值。梯度累積對訓練過程的最后一步進行了修正：在繼續下一個小批之前，保存梯度值，并將新的梯度添加到之前保存的梯度中，用這種方法取代更新每個小批的網絡權重。只有在模型處理了幾個小批次后，才會更新權重。梯度積累模擬了一個更大的批大小，如果想在一個小批中使用64張圖像，如果批大小超過了8，則會報“CUDA內存出錯…”。在這種情況下，可以使用8批圖像，并在模型處理64/8=8批后更新一次權重。如果你從這8個批次中積累每一個梯度，結果將是（幾乎）相同的，這樣便能夠執行訓練啦！
例程：
沒有梯度累積的標準訓練環通常為：
在PyTorch中，梯度累積可以很容易地完成。模型利用accumulation_steps處理完成小批之后，便可以執行優化。還可以利用accumulation_steps根據損失函數的性質來劃分運行損失：
真漂亮，對嗎？當調用loss.backward() 時計算梯度，并由PyTorch累積，直到調用optimizer.zero_grad()時停止。
重點
某些網絡體系結構使用專用的批處理操作，如BatchNorm，當使用相同的批大小時，結果可能會略有不同。
混合精度訓練
概述?
混合精度訓練是指將部分或全部FP32參數轉換為更小的格式，如FP16、TF16（浮點張量）或BF16(浮點字節)。
主要優勢
混合精度訓練的主要優勢是：
- 減少內存使用；
- 性能提速（更高的算術強度或更小的通信占用）；
- 使用專用硬件進行更快地計算。
目前只對第一個優勢感興趣——減少內存的使用量，來看看如何使用PyTorch模型實現它。
例程：
結果，在完成.half()操作之后，模型變小了2倍。將模型轉換為不同的格式(即BF16，TF16)后的縮放損失，將在后續的文章中討論。有些操作在FP16中是無法完成的，如Softmax。PyTorch可利用torch.autocast 來處理這些特殊情況。
8位優化器
增加模型尺寸是獲得更佳性能的有效途徑。然而，訓練大模型時需要存儲模型、梯度和優化器的狀態(例如，Adam的指數平滑和及先前梯度的平方和)，所有這些都存儲在數量有限的可用內存之中。
將32位優化器降到8位優化器，將數值的范圍從232減少到僅2?=256，會對優化器預留的內存數量產生巨大的影響。
研究人員提出了一種新的8位Adam優化器，論文作者在文中這么說： “它將32位的性能維持到部分原始內存中”。
8位優化器有三個組成部分：（1）塊級量化，隔離異常值，將誤差均勻地分配給每一個比特；（2）動態量化，高精度地量化小值和大值；（3）穩定的嵌入層，以提高詞嵌入優化模型的穩定性。
有了這些組件，可直接使用8位狀態執行優化。將8位優化器狀態量化為32位，執行更新，然后再將狀態量化為8位進行存儲。在寄存器中逐元素進行8位到32位的轉換，無需慢速復制到GPU內存或額外的臨時內存中執行量化和去量化。對于GPU來說，這意味著8位優化器要快于常規的32位優化器。
來看看使用8位Adam之后，鼓舞人心的結果：
可以看出，使用量化的Adam可以節省大約8.5 GB的GPU內存，看起來相當棒！
理解了它的可用性之后，再來看看如何用python實現它。
由Facebook提供的Bitsandbytes 包是一個圍繞CUDA自定義函數的輕量級包裝器，封裝了 8位優化器和量化函數，利用它可以實現8位Adam的使用。
例程：
如上所述，量化優化器的使用非常簡單，結果也不錯。
綜合上述全部方法，對GPU上的GPT-2-XL進行微調。
最后，在掌握了上述方法之后，利用這些方法來解決實際問題，對擁有15億個參數的GPT-2-XL模型進行微調。顯然，無法將它加載到12 GB內存的NVIDIA GeForce RTX 3060 GPU之上。列出可以使用的全部方法：
- 梯度檢查點；
- 混合精度訓練(我設了一個技巧：使用相同模型的兩個樣本。首先，用.half將它加載到GPU上，將其命名為gpu_model；其次，在CPU上，將其命名為cpu_model。評估好GPU模型之后，將 gpu_model的梯度加載到cpu_model中，運行optimizer.step()，將更新后的參數加載到gpu_model上)；
- 使用batch_size=64，minibatch_size=4的梯度累積，需要通過 accumulation_steps來縮放損失；
- 8位Adam優化器。
把以上方法全部利用起來，查看一下代碼：
利用上述所有方法之后，在GPU上實現了對16GB的GPT-2-XL模型微調，絕了！
結論
在本博中，給出了高效使用內存的關鍵概念，它適用于多種艱巨的任務，如上文所述。將在后續的文章中討論其他概念。衷心感謝，撥冗閱讀本文!
