
大數(shù)據(jù)文摘授權(quán)轉(zhuǎn)載自數(shù)據(jù)派THU
作者：Stanislav Belyasov 
翻譯：陳之炎 
校對(duì)：趙茹萱

在訓(xùn)練模型過(guò)程中，細(xì)數(shù)那些完勝“CUDA 內(nèi)存出錯(cuò)..”報(bào)錯(cuò)的提高內(nèi)存效率技術(shù)。

提問(wèn)：模型大小超過(guò)GPU 容量怎么辦？

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

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

開始吧！

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

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

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

問(wèn)題的核心

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

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

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

假設(shè)已有一塊帶有12 GB內(nèi)存的NVIDIA GeForce RTX 3060。首先， 1e9個(gè)FP32參數(shù)約占4 GB的GPU內(nèi)存。同樣，對(duì)于梯度，也將保留相同數(shù)量的內(nèi)存。所以，總共已經(jīng)保留了8 GB的內(nèi)存，由于還沒有開始訓(xùn)練，也沒有加載優(yōu)化器，加載優(yōu)化器也同樣需要一定數(shù)量的內(nèi)存。Adam優(yōu)化器需要為每個(gè)參數(shù)存儲(chǔ)第一備份和第二備份，即需要8 GB額外內(nèi)存。

算下來(lái)，必須有大約16 GB的GPU內(nèi)存，才能正確地將模型加載到GPU上，在本文的例子中，GPU只有12 GB的空閑內(nèi)存?？雌饋?lái)很不妙，對(duì)吧？

然而，可以通過(guò)一些方法來(lái)嘗試解決這個(gè)問(wèn)題，以下是相關(guān)內(nèi)容：

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

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

開始

提問(wèn)：模型比GPU容量大，怎么辦？

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

概述

如果模型大于GPU容量，即便將批大小設(shè)為1都不夠，那該怎么辦呢？有一個(gè)解決方案，即設(shè)置梯度檢查點(diǎn)，下面來(lái)看看這個(gè)概念。

對(duì)于一個(gè)簡(jiǎn)單的包含n層的前饋神經(jīng)網(wǎng)絡(luò)來(lái)說(shuō)，梯度的計(jì)算圖如下：

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

下面是這些節(jié)點(diǎn)的計(jì)算順序，紫色陰影圓圈表示在給定時(shí)間里需要將哪個(gè)節(jié)點(diǎn)保存到內(nèi)存之中。

梯度檢查點(diǎn)

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

這種策略在內(nèi)存方面是最優(yōu)的。但是，請(qǐng)注意，節(jié)點(diǎn)計(jì)算的數(shù)量進(jìn)行了n2次縮放，而先前的縮放系數(shù)為n：每個(gè)n個(gè)節(jié)點(diǎn)都按n次順序重新計(jì)算。由于計(jì)算速度較慢，這種方法并不適用于深度學(xué)習(xí)。

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

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

例程：

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

梯度累積/微批次

概述

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

什么是梯度累積？

在訓(xùn)練神經(jīng)網(wǎng)絡(luò)時(shí)，通常會(huì)將數(shù)據(jù)分批量處理，神經(jīng)網(wǎng)絡(luò)預(yù)測(cè)批處理標(biāo)簽，用于計(jì)算相對(duì)于實(shí)際目標(biāo)的損失。接下來(lái)，執(zhí)行反向傳遞計(jì)算出梯度，更新模型權(quán)值。

梯度累積對(duì)訓(xùn)練過(guò)程的最后一步進(jìn)行了修正：在繼續(xù)下一個(gè)小批之前，保存梯度值，并將新的梯度添加到之前保存的梯度中，用這種方法取代更新每個(gè)小批的網(wǎng)絡(luò)權(quán)重。只有在模型處理了幾個(gè)小批次后，才會(huì)更新權(quán)重。

梯度積累模擬了一個(gè)更大的批大小，如果想在一個(gè)小批中使用64張圖像，如果批大小超過(guò)了8，則會(huì)報(bào)“CUDA內(nèi)存出錯(cuò)…”。在這種情況下，可以使用8批圖像，并在模型處理64/8=8批后更新一次權(quán)重。如果你從這8個(gè)批次中積累每一個(gè)梯度，結(jié)果將是（幾乎）相同的，這樣便能夠執(zhí)行訓(xùn)練啦！

例程：

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

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

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

重點(diǎn)

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

混合精度訓(xùn)練

概述

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

主要優(yōu)勢(shì)

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

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

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

例程：

結(jié)果，在完成.half()操作之后，模型變小了2倍。

將模型轉(zhuǎn)換為不同的格式(即BF16，TF16)后的縮放損失，將在后續(xù)的文章中討論。

有些操作在FP16中是無(wú)法完成的，如Softmax。PyTorch可利用torch.autocast 來(lái)處理這些特殊情況。

8位優(yōu)化器

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

將32位優(yōu)化器降到8位優(yōu)化器，將數(shù)值的范圍從232減少到僅2?=256，會(huì)對(duì)優(yōu)化器預(yù)留的內(nèi)存數(shù)量產(chǎn)生巨大的影響。

研究人員提出了一種新的8位Adam優(yōu)化器，論文作者在文中這么說(shuō)： “它將32位的性能維持到部分原始內(nèi)存中”。

8位優(yōu)化器有三個(gè)組成部分：（1）塊級(jí)量化，隔離異常值，將誤差均勻地分配給每一個(gè)比特；（2）動(dòng)態(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位進(jìn)行存儲(chǔ)。在寄存器中逐元素進(jìn)行8位到32位的轉(zhuǎn)換，無(wú)需慢速?gòu)?fù)制到GPU內(nèi)存或額外的臨時(shí)內(nèi)存中執(zhí)行量化和去量化。對(duì)于GPU來(lái)說(shuō)，這意味著8位優(yōu)化器要快于常規(guī)的32位優(yōu)化器。

來(lái)看看使用8位Adam之后，鼓舞人心的結(jié)果：

可以看出，使用量化的Adam可以節(jié)省大約8.5 GB的GPU內(nèi)存，看起來(lái)相當(dāng)棒！

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

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

例程：

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

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

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

列出可以使用的全部方法：

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

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

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

結(jié)論

在本博中，給出了高效使用內(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

點(diǎn)「在看」的人都變好看了哦！

如何在GPU資源受限情況下微調(diào)超大模型

概述

梯度檢查點(diǎn)

例程：

概述

什么是梯度累積？

例程：

重點(diǎn)

概述

主要優(yōu)勢(shì)

例程：

例程：

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

什么是梯度累積？

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