Memory Optimization for Deep Networks

Deep learning is slowly , but steadily , hitting a memory bottleneck . While the tensor computation in top-of-the-line GPUs increased by 32× over the last five years , the total available memory only grew by 2.5× . This prevents researchers from exploring larger architectures , as training large networks requires more memory for storing intermediate outputs . In this paper , we present MONET , an automatic framework that minimizes both the memory footprint and computational overhead of deep networks . MONET jointly optimizes the checkpointing schedule and the implementation of various operators . MONET is able to outperform all prior handtuned operations as well as automated checkpointing . MONET reduces the overall memory requirement by 3× for various PyTorch models , with a 9-16 % overhead in computation . For the same computation cost , MONET requires 1.2-1.8× less memory than current state-of-the-art automated checkpointing frameworks . Our code is available at https : //github.com/utsaslab/MONeT . 1 INTRODUCTION . Deep networks are widely used in domains ranging from image classification ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2015 ; He et al. , 2016 ) to video recognition ( Wu et al. , 2019 ; Feichtenhofer et al. , 2019 ) or natural language processing ( Devlin et al. , 2019 ; Yang et al. , 2019 ) . However , training deep networks is resource-intensive . In particular , the amount of GPU memory bottlenecks training many deep networks ( Dong et al. , 2016 ; Kim et al. , 2016 ; Chen et al. , 2018 ; Child et al. , 2019 ) . This bottleneck requires either modifying the network architecture or scaling training to multiple nodes , incurring significant overheads . We presents MONET , an automatic framework to minimize memory footprint for deep networks . MONET jointly optimizes global compute-graph-level techniques ( such as checkpointing ) and local techniques ( such as memory-efficient implementations of individual operator ) . At the heart of MONET is a theoretical analysis that enables joint optimization and provides tight bounds on memory consumption . We analyze the memory consumption and computational cost of a general forward and backward pass under changing local operator implementations and a global checkpointing schedule . Specifically , we are able to tightly bound the peak memory consumption for network forward , backward , and recomputation stages . MONET uses these constraints to optimize for the most efficient forward and backward implementation both locally and globally under a fixed memory budget . We linearize all memory bounds , and express both implementation selection and checkpointing as a 0-1 integer program , which we solve using standard solvers . We conduct extensive experiments , demonstrating that MONET significantly outperforms existing automatic frameworks that use local or global techniques . On multiple architectures ( ResNet ( He et al. , 2016 ) , VGG ( Simonyan & Zisserman , 2015 ) , UNet ( Ronneberger et al. , 2015 ) , GoogleNet ( Szegedy et al. , 2015 ) , MobileNet-V2 ( Sandler et al. , 2018 ) ) , memory budgets ( 5-10 GB ) , and network configurations ( multiple resolutions ) , MONET consistently achieves lower memory footprints at equivalent or lower computational overhead . MONET reduces the overall memory requirement by 3× for various models , with a 9-16 % overhead in computation . For the same computation cost , MONET requires 1.2-1.8× less memory than the current state-of-the-art automated checkpointing framework . The results achieved by MONET demonstrate the power of jointly optimizing global checkpointing schedules and local operator implementations . 2 RELATED WORK . There are two broad families of approaches to reduce the memory footprint of a deep network during training : operator-level implementation changes , and global , graph-level optimizations . The novel aspect of MONET is that it is able to combine both approaches and find the optimal mix of local and global techniques for a given network . Operator-Specific Optimizations . Researchers have found creative ways to implement individual operators or groups of operators in a more memory-efficient manner . Standard deep learning frameworks ( Jia et al. , 2014 ; Collobert et al. , 2011 ; Paszke et al. , 2019 ; Abadi et al. , 2016 ) provide different implementations of certain operators that trade computation for intermediate memory use . These implementation are chosen according to local search heuristics , and are not globally optimal . Gist ( Jain et al. , 2018 ) proposes several hand-crafted optimizations such as storing only ReLU signs . RevNets ( Gomez et al. , 2017 ) redesigns a ResNet ( He et al. , 2016 ) architecture making each network block reversible , thereby eliminating the need to store intermediate activations for backpropagation . Memory-efficient DenseNets ( Pleiss et al. , 2017 ) reduce memory utilized for feature maps by recomputing all intermediate feature maps during the backward pass with a small compute overhead . In-place activated batchnorm ( Bulò et al. , 2018 ) or ReLU layers use output activations to compute their gradients , thus reusing a single memory buffer for the gradient computation in consecutive layers . Mixed precision training ( Micikevicius et al. , 2018 ) uses half precision ( FP16 ) instead of single precision ( FP32 ) for all tensors and arithmetic during training , reducing the memory by nearly half . While training at precision lower than FP16 results in loss of training quality ( Banner et al. , 2018 ) , prior work like backpropagation with approximate activations ( Chakrabarti & Moseley , 2019 ) carefully quantize certain intermediate outputs ( activations ) to 4 bits , resulting in significant memory savings . Although these hand-crafted techniques independently result in memory savings , there is no one-size-fits-all recipe , and different implementations perform best on different architectures . In contrast , MONET automatically finds the best implementation for each forward and backward operator given a memory budget . Checkpointing . Chen et al . ( 2016 ) proposed dividing a network into different segments , dropping all intermediate outputs within each segment , and recomputing them later . Chen et al . use √ n equal segments , trading memory savings for the cost of an extra forward pass . Checkmate ( Jain et al. , 2019 ) solves the problem in a more general setting , using an mixed-integer linear program solver to decide which layers to recompute for a given network . Like Checkmate , our work optimizes a checkpointing schedule , but on a different computation graph . Our computation graph allows for the optimization of an entire execution plan jointly finding a checkpointing schedule and the best implementation of each forward and backward operator . In Checkmate , changes in operator implementation induce a different computation graph , and could thus not directly be optimized . Appendix F highlights some of the difficulties of adding operator optimizations into Checkmate . In summary , while much work has been done on local optimizations ( operator implementations ) and global compute-graph-level techniques ( automated checkpointing ) , MONET is the first system to jointly optimize a given architecture using both local and global techniques . 3 PRELIMINARIES . Let the forward pass of a CNN with parameters Θ be expressed as a directed-acyclic graph ( DAG ) , where each node i ∈ { 1 , . . . , N } corresponds to an operator forwardi , and edges ( i , j ) ∈ E specify the data-flow dependencies , i.e. , the output of operator i is used as input in operator j . Without loss of generality , computational dependency ( i , j ) ∈ E implies i < j . Let Nj = { i : ( i , j ) ∈ E } be the set of all incoming edges of an operation j . We will first discuss the forward pass through a network and the basic form of a backward pass using checkpointing . The backward pass reverses all computational dependency expressed in our DAG , and induces certain dependencies on forward activations . We call these checkpoint dependencies Dk . They are either saved or recomputed depending on a schedule ( s , r ) . Checkpointing creates a trade-off between computation and memory consumption . To highlight this trade-off , we formally compute the amount of memory consumed in both forward and backward passes , which allows us to optimize for the ideal execution plan in Sec . 4 . We provide a reference to the notations introduced in this section and the next along with their explanations in Appendix A . The Forward Pass . Alg . 1 shows a general overview of the forward pass in a deep network , as implemented in standard deep learning frameworks ( Jia et al. , 2014 ; Collobert et al. , 2011 ; Paszke et al. , 2019 ; Abadi et al. , 2016 ) . The algorithm proceeds in increasing order of index i . Each operator forwardi ( · ) depends on a set of tensors L stored in local memory . These tensors include model parameters Θ , computational dependenciesNi , and tensors stored for later forward operators , i.e . skip or residual activations ( He et al. , 2016 ) . At each iteration , we add any output tensors of forwardi to the local memory L. Early deep learning frameworks ( Jia et al. , 2014 ; Collobert et al. , 2011 ) strictly grew the set of local tensors L leading to an unnecessarily high memory consumption . Modern graph-based frameworks ( Paszke et al. , 2019 ; Abadi et al. , 2016 ) reduce the memory footprint by aggressively pruning local memory L and freeing any tensor that is no longer used in later computations . Some output activations xi are used in the backward pass , and have to be saved for later . We use a checkpointing schedule sN to determine which . Formally , sNi ∈ { 0 , 1 } indicates whether the activation of node i is stored during the forward pass . An activation which is not stored will be recomputed if it is needed during the backward pass . Analyzing peak memory consumption of the forward pass . Only the forwardi operator ( Alg . 1 L. 4 ) allocates memory . All other operators perform mere bookkeeping on existing tensor . It is thus sufficient to study the peak memory consumption mNi in forwardi for each node i . Let Li , S N i be the set of local tensors L and saved tensors S while calling forwardi respectively . Li includes all parameters and computational dependencies for this and later forward passes Li = Θ ∪ { xj : j ∈ Nt for any t ≥ i and j < i } . Li is constant and computed ahead of time . The schedule sN determines the set of saved tensors SNi = { xj : sNj = 1 for j < i } . In addition , each forward operator uses a certain amount of workspace memory ci to store intermediate results . The total memory consumption of a forward operator is thus mi = ci + |xi|+ |SNi ∪ Li| = ci + |xi|+ ∑ xj∈Li |xj |+ ∑ j < i : xj /∈Li |xj |sNj , ( 1 ) where | · | refers to the memory consumed by a tensor or set of tensors . Most of the memory consumption is constant and does not depend on the schedule . The Backward Pass . The backward pass proceeds in a reverse order , as summarized in Alg . 2. backwardk ( · ) of each node k depends on a set of gradient tensors L̂ and forward tensors { xi : i ∈ Dk } . Any gradients required by the current and later backward passes are stored in local memory L̂ . Dependencies Dk may either be stored in Sk or need to be recomputed from checkpoints in Sk . Recomputation involves forward computation of one or more nodes , which increases computational overhead , and allows for a new set of tensors Sk−1 to be saved . After recomputation , all dependencies Dk are kept in memory . The backward operation produces a gradient for each input tensor of the original forward operation , which is added to L̂ if required for a later backward computation . We aggressively remove tensors in L̂ that are not required . Analyzing the peak memory consumption of the backward pass . Peak memory consumption m̂k again only depends on the forwardi ( Alg . 2 L. 7 ) and backwardk ( Alg . 2 L. 12 ) operations . For the backwardk operation , let ĉk be the workspace memory , L̂k be the set of gradient tensors stored , Dk = { xi : i ∈ Dk } be the forward tensors used , and Sk−1 be the set of newly saved tensors . Here L̂k and Dk can be pre-computed . The total memory consumption for the backwardk call is m̂k = ĉk + |yk|+ |Sk−1 ∪ L̂k ∪Dk| = ĉk + |yk|+ ∑ yl∈L̂k |yl|+ ∑ xi∈Dk |xi|+ ∑ xi /∈Dk sk−1i |xi| . ( 2 ) Here again , only the last term depends on the checkpointing schedule , while the rest is a constant . Analyzing the peak memory consumption of the recomputation . Finally , the peak memory m̃ki for the forwardi call ( Alg . 2 L. 7 ) depends on the set of local tensors L , checkpoint dependencies D , saved tensors S , and gradient tensors L̂ , named Lki , Dk , S k−1 i , L̂k respectively . Following the forward pass : m̃ki = ci + |xi|+ |L̂k|+ |Sk−1i ∪ L k i ∪Dk| = ci + |xi|+ |L̂k|+ ∑ j < i : xj /∈Lki ∪Dk sk−1j |xj |+ ∑ j < i : xj∈Lki ∪Dk |xj |+ ∑ j > i skj |xj | . ( 3 ) Unlike the forward pass , Lki is no longer constant , but depends on past saved tensors and future recomputations in ( s , r ) : Lki = Θ ∪ { xj : j ∈ Nt for any t ≥ i with rkt = 1 and j < i } . In the next section , we show how to take this formalization of the forward and backward pass , and find an optimal execution plan including checkpointing schedule ( s , r ) , forwardi implementations , and backwardk implementations , under a fixed memory budget .

Summary: Training deep learning models is becoming increasingly challenging due to a memory bottleneck that limits the size of the feature maps that can be stored. The paper presents an automatic framework (MONET) that minimizes the memory footprint for deep networks. The novelty of MONET is that it jointly optimizes over: (a) global compute graph level techniques (such as checkpointing) and (b) local techniques (such as memory-efficient implementations of individual operators). While there are several existing works that focus separately on optimizing global techniques (e.g. the work on “Checkmate”) or local techniques, MONET is the first to jointly optimize over global and local techniques.