GMLake：基于虚拟内存拼接的大规模 DNN 训练高效透明 GPU 内存去碎片化
分布式、并行与集群计算
2024-01-17 v1
摘要
大规模深度神经网络（DNN），如大型语言模型（LLM），彻底改变了人工智能（AI）领域并日益普及。然而，训练或微调此类模型需要巨大的计算能力和资源，其中单个加速设备（如 GPU）的内存容量是最主要的瓶颈之一。由于 GPU 原生内存分配器的开销过高（例如 ），PyTorch 和 TensorFlow 等 DNN 框架采用了缓存分配器，该分配器维护一个具有分裂机制的内存池以实现快速的内存（分）配。不幸的是，对于重计算、卸载、分布式训练和低秩适配等流行的内存缩减技术，缓存分配器的效率会迅速下降。主要原因是这些内存缩减技术引入了频繁且不规则的内存（分）配请求，导致基于分裂的缓存分配器出现严重的碎片化问题。为了缓解这一碎片化问题，我们提出了一种基于底层 GPU 虚拟内存管理的新型内存分配框架，称为 GPU 内存湖（GMLake）。GMLake 采用了一种新颖的虚拟内存拼接（VMS）机制，该机制可以通过虚拟内存地址映射融合或合并非连续的内存块。在具有 80 GB 内存的 GPU A100 上，针对八个 LLM 模型，GMLake 平均可减少 9.2 GB（最高达 25 GB）的 GPU 内存使用和 15%（最高达 33%）的碎片化。GMLake 对 DNN 模型和内存缩减技术完全透明，确保了资源密集型深度学习任务的无缝执行。我们已在 https://github.com/intelligent-machine-learning/glake/tree/main/GMLake 开源了 GMLake。
引用
@article{arxiv.2401.08156,
title = {GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching},
author = {Cong Guo and Rui Zhang and Jiale Xu and Jingwen Leng and Zihan Liu and Ziyu Huang and Minyi Guo and Hao Wu and Shouren Zhao and Junping Zhao and Ke Zhang},
journal= {arXiv preprint arXiv:2401.08156},
year = {2024}
}
备注
Accepted by ASPLOS24