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

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

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

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

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

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

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