一句话：提出基于 GPU 虚拟内存拼接的分配框架 GMLake，透明缓解 DNN 训练中的内存碎片并降低显存占用。
大规模 DNN（如 LLM）训练与微调受限于单 GPU 显存容量。原生分配器开销高，PyTorch/TensorFlow 等采用缓存分配器与分裂机制，但在重计算、卸载、分布式训练、低秩适配等内存优化下，频繁且不规则的分配/释放导致严重碎片，降低分配效率。GMLake 是一个基于底层 GPU 虚拟内存管理的内存分配框架，其核心是虚拟内存拼接（virtual memory stitching, VMS）：通过虚拟地址映射融合或组合非连续内存块，无需移动物理数据即可缓解碎片。GMLake 对 DNN 模型和内存优化技术完全透明。在 A100 80GB 上评估八个 LLM 模型：平均减少 9.2GB（最高 25GB）显存占用，碎片降低 15%（最高 33%）。
Large-scale deep neural networks (DNNs), such as large language models (LLMs), have revolutionized the artificial intelligence (AI) field and become increasingly popular. However, training or fine-tuning such models requires substantial computational power and resources, where the memory capacity of a single acceleration device like a GPU is one of the most important bottlenecks. Owing to the prohibitively large overhead (e.g., 10×) of GPUs' native memory allocator, DNN frameworks like PyTorch and TensorFlow adopt a caching allocator that maintains a memory pool with a splitting mechanism for fast memory (de)allocation. Unfortunately, the caching allocator's efficiency degrades quickly for popular memory reduction techniques such as re-computation, offloading, distributed training, and low-rank adaptation. The primary reason is that those memory reduction techniques introduce frequent and irregular memory (de)allocation requests, leading to severe fragmentation problems for the splitting-based caching allocator. To mitigate this fragmentation problem, we propose a novel memory allocation framework based on low-level GPU virtual memory management called GPU memory lake (GMLake). GMLake employs a novel virtual memory stitching (VMS) mechanism, which can fuse or combine non-contiguous memory blocks with a virtual memory address mapping. GMLake can reduce average of 9.2 GB (up to 25 GB) GPU memory usage and 15% (up to 33%) fragmentation among eight LLM models on GPU A100 with 80 GB memory. GMLake is completely transparent to the DNN models and memory reduction techniques and ensures the seamless execution of resource-intensive deep-learning tasks. We have open-sourced GMLake at https://github.com/intelligent-machine-learning/glake/tree/main/GMLake.
DOI 原文 · arXiv · PDF（开放获取） · DBLP
@inproceedings{0003ZXLLHGWZZZ24,
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},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620665.3640423}
}