Training Feature-Awared GPU-Memory Allocation and Management for Deep Neural Networks
摘要
Memory limitation of a single GPU is an urgent problem for efficient training of deep learning models on multi-GPU clusters, as a large number of model parameters, intermediate states, and activation values need to be stored during training. Efficient management and allocation of GPU memory is the key to improving memory resource utilization and model training efficiency. The existing methods mainly optimize the layout of memory space when training data is allocated and released to reduce memory waste. These methods do not consider the structural differences between models and the memory access feature of different data (model parameters, intermediate data, etc.) during the model training, so there are still problems of memory fragmentation and low utilization. To address these problems, this paper proposes an efficient GPU-memory allocation and management method, TMManager. It proposes an access feature analyzer, sampling and analyzing the structure of the model and the access features of memory for different data in the model training. Then, we design a dual-level memory partition management method with block and chunk, and a time-sharing deque memory allocation method, to reduce memory fragmentation and improve memory utilization. The experiments demonstrate that TMManager can save up to 23.5%, 59.9% of memory space compared with the memory allocators of TensorFlow, and Pytorch. Compared to PagedAttention, TMManager also realizes a faster and more convenient way of memory allocation.