AI訓(xùn)練的最大障礙不是算力，而是“內(nèi)存墻”

子豪 2021-04-01 16:08:06 來(lái)源：量子位

子豪 發(fā)自 凹非寺
量子位 報(bào)道 | 公眾號(hào) QbitAI

AI訓(xùn)練的計(jì)算量每年都在大幅增長(zhǎng)，最近有研究指出，AI訓(xùn)練未來(lái)的瓶頸不是算力，而是GPU內(nèi)存。

AI加速器通常會(huì)簡(jiǎn)化或刪除其他部分，以提高硬件的峰值計(jì)算能力，但是卻難以解決在內(nèi)存和通信上的難題。

無(wú)論是芯片內(nèi)部、芯片間，還是AI加速器之間的通信，都已成為AI訓(xùn)練的瓶頸。

Transformer模型中的參數(shù)數(shù)量（紅色）呈現(xiàn)出2年240倍的超指數(shù)增長(zhǎng)，而單個(gè)GPU內(nèi)存（綠色）僅以每2年2倍的速度擴(kuò)大。

△多年來(lái)SOTA模型的參數(shù)數(shù)量（紅點(diǎn)）以及AI加速器存儲(chǔ)容量（綠點(diǎn)）的演變

訓(xùn)練AI模型的內(nèi)存需求，通常是參數(shù)數(shù)量的幾倍。因?yàn)橛?xùn)練需要存儲(chǔ)中間激活，通常會(huì)比參數(shù)（不含嵌入）數(shù)量增加3-4倍的內(nèi)存。

于是，AI訓(xùn)練不可避免地撞上了“內(nèi)存墻”（Memory Wall），內(nèi)存墻不僅是內(nèi)存容量，也包括內(nèi)存?zhèn)鬏攷挕?/p>

在很多情況下，數(shù)據(jù)傳輸?shù)娜萘亢退俣?，都沒(méi)有觸摸到內(nèi)存墻。

△訓(xùn)練不同神經(jīng)網(wǎng)絡(luò)模型所需的內(nèi)存量

從圖中可以看出，每當(dāng)GPU內(nèi)存容量增加時(shí)，開(kāi)發(fā)人員就會(huì)設(shè)計(jì)出新模型；

2019年GPT-2所需的內(nèi)存容量，已經(jīng)是2012年的AlexNet的7倍以上；

自谷歌團(tuán)隊(duì)在2017年提出Transformer，模型所需的內(nèi)存容量開(kāi)始大幅增長(zhǎng)。

為什么不能靠多GPU堆顯存

那么，為了擺脫單個(gè)硬件的有限內(nèi)存容量和帶寬，是否可以將訓(xùn)練擴(kuò)展到多個(gè)AI加速器，使用分布式內(nèi)存呢？

事實(shí)上，這樣做也會(huì)面臨內(nèi)存墻的問(wèn)題，神經(jīng)網(wǎng)絡(luò)加速器之間移動(dòng)數(shù)據(jù)的通信瓶頸，甚至比芯片上的數(shù)據(jù)移動(dòng)還慢且低效。

與單系統(tǒng)內(nèi)存的情況類似，擴(kuò)展帶寬的技術(shù)難題還尚未被攻克。僅在很少的通信和數(shù)據(jù)傳輸?shù)那闆r下，橫向擴(kuò)展才適用于計(jì)算密集型問(wèn)題。

從圖中可以看出，20年間，運(yùn)算設(shè)備的算力提高了90,000倍；

雖然存儲(chǔ)器從DDR發(fā)展到GDDR6x，能夠用于顯卡、游戲終端和高性能運(yùn)算，接口標(biāo)準(zhǔn)也從PCIe1.0a升級(jí)到NVLink3.0；

但是和算力的提高幅度相比，通訊帶寬的增長(zhǎng)只有30倍，可以說(shuō)非常緩慢。

由于算力和內(nèi)存之間的差距越來(lái)越大，訓(xùn)練更大的模型也會(huì)更加困難。

怎樣突破“內(nèi)存墻”

怎樣解決內(nèi)存限制問(wèn)題？作者從三個(gè)方面進(jìn)行了分析。

訓(xùn)練算法的改進(jìn)

訓(xùn)練神經(jīng)網(wǎng)絡(luò)模型的一大挑戰(zhàn)，就是要進(jìn)行蠻力超參數(shù)調(diào)整。雖然可以通過(guò)二階隨機(jī)優(yōu)化方法來(lái)實(shí)現(xiàn)，不過(guò)目前的方法卻增加了3-4倍的內(nèi)存占用，這一點(diǎn)仍需解決。

微軟的Zero方法（一種萬(wàn)億級(jí)模型參數(shù)訓(xùn)練方法），實(shí)現(xiàn)了在相同內(nèi)存下，通過(guò)去除多余的優(yōu)化狀態(tài)變量，來(lái)訓(xùn)練8倍大的模型。

也可以在傳遞過(guò)程中只存儲(chǔ)或檢查激活的子集，而不保存所有激活，以此將內(nèi)存占用減少5倍，不過(guò)需要增加20%的計(jì)算量。

此外，從單精度算法到半精度（FP16）算法的進(jìn)展，使硬件計(jì)算能力提高了10倍以上，可以進(jìn)一步研究適合INT8精讀訓(xùn)練的優(yōu)化算法。

高效部署

最新的SOTA模型（例如：GPT-3）需要分布式內(nèi)存部署，這是一個(gè)很大的挑戰(zhàn)?？梢酝ㄟ^(guò)降低精度或刪除其冗余參數(shù)，來(lái)壓縮這些模型，以進(jìn)行推理。

在訓(xùn)練或推理過(guò)程中，可以降低至INT4精度，模型占用空間和延遲能夠減少8倍。不過(guò)，想要將訓(xùn)練精度降低到遠(yuǎn)低于FP16，仍然很困難。

而刪除冗余參數(shù)，則可能導(dǎo)致準(zhǔn)確率下降。當(dāng)前的方法能夠修剪30％的具有結(jié)構(gòu)化稀疏性的神經(jīng)元，以及80％的具有非結(jié)構(gòu)化稀疏性的神經(jīng)元，以保證對(duì)準(zhǔn)確性的影響最小。

AI加速器的設(shè)計(jì)

雖然很難同時(shí)提高存儲(chǔ)帶寬和峰值計(jì)算能力，但是可以犧牲峰值計(jì)算，以獲得更好的帶寬。

在帶寬受限問(wèn)題上，CPU的性能要比GPU好得多，但是與相比GPU相比，CPU的峰值計(jì)算能力要小一個(gè)數(shù)量級(jí)左右。

因此，可以研究一種在二者之間的另一種架構(gòu)，實(shí)現(xiàn)更高效的緩存。

研究數(shù)據(jù)可戳下方鏈接查看~

原文鏈接：

https://medium.com/riselab/ai-and-memory-wall-2cb4265cb0b8
https://github.com/amirgholami/ai_and_memory_wall