訓(xùn)練大模型時(shí)，顯存都哪去了？ 原創(chuàng)
GPT-2（XL）有15億個(gè)參數(shù)，使用16位精度，一個(gè)參數(shù)占用2個(gè)字節(jié)的內(nèi)存，因此這些參數(shù)大約占用3GB的內(nèi)存。
按照如下超參數(shù)設(shè)置：
- 優(yōu)化器 → Adam
- 批量大小 → 32
- 變換層數(shù)量 → 48
- 序列長(zhǎng)度 → 1000
要想在單個(gè)GPU上訓(xùn)練GPT-2，所需的最小內(nèi)存大概是多少？
答案可能會(huì)嚇到你。
在一個(gè)擁有32GB內(nèi)存的單個(gè)GPU上，幾乎無(wú)法訓(xùn)練一個(gè)3GB的GPT-2模型。
但這怎么可能呢？?jī)?nèi)存都去哪了？讓我們來(lái)了解一下。
模型在訓(xùn)練過(guò)程中有很多方面會(huì)持續(xù)占用內(nèi)存。
#1)優(yōu)化器狀態(tài)，梯度，模型參數(shù)
混合精度訓(xùn)練廣泛用于加速模型訓(xùn)練。
顧名思義，這個(gè)方法的思想是在訓(xùn)練過(guò)程中同時(shí)使用float16低精度（在卷積和矩陣乘法等操作中）和高精度（如32位浮點(diǎn)數(shù)，float32）。
這就是“混合精度”名稱的由來(lái)。
前向傳播和反向傳播都使用16位浮點(diǎn)數(shù)表示權(quán)重和梯度。
因此，如果模型有Φ個(gè)參數(shù)，那么：
● 權(quán)重將占用2 * Φ字節(jié)的內(nèi)存。
● 梯度將占用2 * Φ字節(jié)的內(nèi)存。
這里的“2”表示每個(gè)參數(shù)占用2個(gè)字節(jié)的內(nèi)存（16位）。
Adam 是最受歡迎的模型訓(xùn)練優(yōu)化器之一。
雖然許多實(shí)踐者僅僅因?yàn)樗餍卸褂盟?，但他們沒(méi)有意識(shí)到，在訓(xùn)練過(guò)程中，Adam 會(huì)存儲(chǔ)兩種優(yōu)化器狀態(tài)來(lái)計(jì)算更新——梯度的動(dòng)量和方差。
因此，如果模型有Φ個(gè)參數(shù)，那么這兩個(gè)優(yōu)化器狀態(tài)將消耗：
● 4 * Φ 字節(jié)用于動(dòng)量。
● 另需 4 * Φ 字節(jié)用于方差。
這里的“4”表示每個(gè)參數(shù)占用 4 個(gè)字節(jié)的內(nèi)存（32 位）。
此外，反向傳播結(jié)束時(shí)的更新仍然在32位精度下進(jìn)行，以確保有效的計(jì)算。這導(dǎo)致：
● 另需 4 * Φ 字節(jié)用于模型參數(shù)。
讓我們把它們加起來(lái)：
這就是 16 * Φ，或者 24GB 的內(nèi)存，遠(yuǎn)遠(yuǎn)高于 16 位參數(shù)所使用的 3GB 內(nèi)存。
而且我們還沒(méi)有考慮到所有的因素。
2#)激活值
對(duì)于像大型深度學(xué)習(xí)模型（如大語(yǔ)言模型，LLMs）來(lái)說(shuō)，激活值在訓(xùn)練過(guò)程中占用了大量?jī)?nèi)存。
更確切地說(shuō)，在GPT-2的一個(gè)Transformer塊中計(jì)算的激活值總數(shù)是：
因此，在所有的Transformer塊中，總計(jì)就是：
這是 GPT-2-XL 的配置：
總共大約是 300 億個(gè)激活值。由于每個(gè)激活值使用 16 位表示，所有激活值總共占用 60GB 的內(nèi)存。
通過(guò)使用像梯度檢查點(diǎn)（在上一章討論過(guò)的）這樣的技術(shù)，可以將內(nèi)存消耗降低到大約 8-9GB，但這也會(huì)帶來(lái)額外 25-30% 的計(jì)算開(kāi)銷(xiāo)。
除了可以計(jì)算的內(nèi)存開(kāi)銷(xiāo)外，還有一些額外的開(kāi)銷(xiāo)，例如內(nèi)存碎片化。
內(nèi)存碎片化是指在分配的內(nèi)存塊之間存在小的未使用間隙，導(dǎo)致可用內(nèi)存的低效使用。
內(nèi)存分配請(qǐng)求失敗是因?yàn)闆](méi)有足夠的連續(xù)內(nèi)存塊可用。
在上述討論中，我們考慮了一個(gè)相對(duì)較小的模型——GPT-2（XL），它有 15 億個(gè)參數(shù)，與如今訓(xùn)練的模型規(guī)模相比非常小。
然而，這個(gè)討論可能幫助你反思構(gòu)建大規(guī)模語(yǔ)言模型（LLMs）時(shí)的固有挑戰(zhàn)。很多人常說(shuō)，GPT 模型只是簡(jiǎn)單地堆疊更多的層并使網(wǎng)絡(luò)變得更大。
如果真是那么簡(jiǎn)單，大家都會(huì)在做了。從這個(gè)討論中，你可能已經(jīng)理解到，這并不像僅僅添加更多層那么簡(jiǎn)單。
即便是增加一層，也可能導(dǎo)致額外數(shù) GB 的內(nèi)存需求。多 GPU 訓(xùn)練是這些模型的核心技術(shù)，我們將在另一篇文章中討論。
本文轉(zhuǎn)載自公眾號(hào)人工智能大講堂
原文鏈接：??https://mp.weixin.qq.com/s/PFQZnqJJ-tjFcv6oSjaV-A??
