Capuchin: Tensor-based GPU Memory Management for Deep Learning

By Xuan Peng, Xuanhua Shi, Hulin Dai, Hai Jin, Weiliang Ma, Qian Xiong, Fan Yang, Xuehai Qian. 2020

205 citations recorded by the source corpus.

In recent years, deep learning has gained unprecedented success in various domains, the key of the success is the larger and deeper deep neural networks (DNNs) that achieved very high accuracy. On the other side, since GPU global memory is a scarce resource, large models also pose a significant challenge due to memory requirement in the training process. This restriction limits the DNN architecture exploration flexibility. In this paper, we propose Capuchin, a tensor-based GPU memory management module that reduces the memory footprint via tensor eviction/prefetching and recomputation. The key feature of Capuchin is that it makes memory management decisions based on dynamic tensor access pattern tracked at runtime. This design is motivated by the observation that the access pattern to tensors is regular during training iterations. Based on the identified patterns, one can exploit the total memory optimization space and offer the fine-grain and flexible control of when and how to perform memory optimization techniques. We deploy Capuchin in a widely-used deep learning framework, Tensorflow, and show that Capuchin can reduce the memory footprint by up to 85% among 6 state-of-the-art DNNs compared to the original Tensorflow. Especially, for the NLP task BERT, the maximum batch size that Capuchin can outperforms Tensorflow and gradient-checkpointing by 7x and 2.1x, respectively. We also show that Capuchin outperforms vDNN and gradient-checkpointing by up to 286% and 55% under the same memory oversubscription.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, ML Systems & Hardware

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In ML Systems & Hardware

Citations in the corpus

vPipe: A Virtualized Acceleration System for Achieving Efficient and Scalable Pipeline Parallel DNN Training (2022, methodology): “In addition, an alternative approach to apply dynamic swap and recompute systems (i.e., Capuchin [37]) to distributed settings is to integrate Capuchin to each worker of data parallelism.”
OLLA: Optimizing the Lifetime and Location of Arrays to Reduce the Memory Usage of Neural Networks (2022, background): “This has been demonstrated to effectively offload the data stored on a GPU device onto the host memory [64, 38, 54], but again increases training time due to extra memory transfers.”
Sentinel: Efficient Tensor Migration and Allocation on Heterogeneous Memory Systems for Deep Learning (2021, methodology): “To manipulate tensor allocation, Sentinel-GPU intercepts TensorFlow GPU memory allocators (such as Al-locateRaw and gpu_bfc_Allocator), similar to [9], Sentinel-GPU replaces those allocators with the customized ones for pinned memory control or tensor collocation. add_layer() is implemented as a…”
Parallel Training of Pre-Trained Models via Chunk-Based Dynamic Memory Management (2021, methodology): “Work [29] also combines offloading and checkpointing techniques based on dynamic tensor access pattern tracked at runtime.”
POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging (2022, methodology): “Capuchin (Peng et al., 2020) uses the Memory Saving Per Second (MSPS) heuristic to decide what to page.”
Zico: Efficient GPU Memory Sharing for Concurrent DNN Training (2021, methodology): “, following topological sort, thus keeping the corresponding memory operations ordered across iterations [32].”