近日，從螞蟻集團(tuán)獲悉，螞蟻集團(tuán)和上海交通大學(xué)合作的技術(shù)成果GMLake被計(jì)算機(jī)體系結(jié)構(gòu)四大頂級(jí)會(huì)議之一的ASPLOS 24接收。

據(jù)悉，這篇名為《GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching》的研究成果，針對(duì)業(yè)界普遍存在的大模型訓(xùn)練顯存效率問題，提出了一種高效的動(dòng)態(tài)顯存組合技術(shù)VMS（Virtual Memory Stitching），最高提高33%的GPU可用顯存，使顯存使用更為高效，進(jìn)而緩解了“內(nèi)存墻”問題。ASPLOS會(huì)議曾推動(dòng)了多核處理器、深度學(xué)習(xí)處理器等計(jì)算機(jī)系統(tǒng)領(lǐng)域的核心創(chuàng)新，在學(xué)術(shù)界和工業(yè)界擁有顯著的影響力。

以ChatGPT為代表的生成式大模型正在引領(lǐng)當(dāng)前AI發(fā)展。為了訓(xùn)練得到如此高效的大模型，需要強(qiáng)大和昂貴的基礎(chǔ)算力支撐。而當(dāng)前制約高效訓(xùn)練的因素中，除了算力，顯存容量也非常關(guān)鍵，即“內(nèi)存墻”問題。

業(yè)界已有很多優(yōu)化工作，包括模型、框架層的改造，甚至為之犧牲模型精度。且業(yè)界廣泛使用的訓(xùn)練框架如PyTorch存在顯存管理效率不高，顯存資源碎片化的現(xiàn)象，即顯存中出現(xiàn)了很多的空洞，使得零散的顯存資源難以有效利用起來，特別是當(dāng)使用了業(yè)界多種優(yōu)化手段之后，碎片化問題反而更為突出，使得本就緊缺的顯存資源浪費(fèi)嚴(yán)重。

對(duì)此，螞蟻集團(tuán)和上海交通大學(xué)提出了一種高效的動(dòng)態(tài)顯存組合技術(shù)框架，稱為GMLake。該框架采用虛擬、物理兩層指針，將碎片化的顯存靈活地拼接起來，從而實(shí)現(xiàn)了聚零為整。GMLake對(duì)典型大模型如GPT、GLM、Vicuna等進(jìn)行了詳細(xì)評(píng)測(cè)，最高擴(kuò)大了約33%的可用顯存，即在80GB的A100 GPU上節(jié)省約25GB的顯存，訓(xùn)練吞吐提高最多4倍。

據(jù)了解，GMLake的使用非常友好，目前已在PyTorch框架上完成集成，對(duì)上層模型代碼和框架代碼完全透明，也就是說模型不需要修改任何代碼即可使用，真正做到了開箱即用。

目前，GMLake的代碼已經(jīng)開源到了GLake項(xiàng)目中。GLake是螞蟻集團(tuán)自研的GPU顯存+傳輸優(yōu)化開源項(xiàng)目，在訓(xùn)練、推理框架底層以及系統(tǒng)層對(duì)顯存和傳輸進(jìn)行一體優(yōu)化，更好應(yīng)對(duì)“內(nèi)存墻”和“傳輸墻”挑戰(zhàn)。