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

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

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

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

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

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

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