大幅提升GPU可用显存量，蚂蚁、上海交大技术成果GMLake入选国际顶会

近日，从蚂蚁集团获悉，蚂蚁集团和上海交通大学合作的技术成果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显存+传输优化开源项目，在训练、推理框架底层以及系统层对显存和传输进行一体优化，更好应对“内存墙”和“传输墙”挑战。