提高 33% 可用显存！蚂蚁 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 显存+传输优化开源项目，在训练、推理框架底层以及系统层对显存和传输进行一体优化，更好应对「内存墙」和「传输墙」挑战。
