M EMORY O PTIMIZATION FOR D EEP N ETWORKS
A BSTRACT
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 ﬁve 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 MON E T , an automatic framework that minimizes both the memory footprint and computational overhead of deep networks.
MON E T jointly optimizes the checkpointing schedule and the implementation of various operators.
MON E T is able to outperform all prior handtuned operations as well as automated checkpointing.
MON E T reduces the overall memory requirement by 3 × for various PyTorch models, with a 9-16 % overhead in computation.
For the same computation cost, MON E T 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 .
I NTRODUCTION
Deep networks are widely used in domains ranging from image classiﬁcation (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 signiﬁcant overheads.
We presents MON E T , an automatic framework to minimize memory footprint for deep networks.
MON E T jointly optimizes global compute-graph-level techniques (such as checkpointing) and local techniques (such as memory-efﬁcient implementations of individual operator).
At the heart of
MON E T 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.
Speciﬁcally, we are able to tightly bound the peak memory consumption for network forward, backward, and recomputation stages.
MON E T uses these constraints to optimize for the most efﬁcient forward and backward implementation both locally and globally under a ﬁxed 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 MON E T signiﬁcantly 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 (5GB), and network conﬁgurations (multiple resolutions), MON E T consistently achieves lower memory footprints at equivalent or lower computational overhead.
MON E T reduces the overall memory requirement by 3 × for various models, with a 9-16 % overhead in computation.
For the same computation cost, MON E T requires 1.2-1.8 × less memory than the current state-of-the-art automated checkpointing framework.
The results achieved by MON E T demonstrate the power of jointly optimizing global checkpointing schedules and local operator implementations.
R ELATED W ORK
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 MON E T is that it is able to combine both approaches and ﬁnd the optimal mix of local and global techniques for a given network.
Operator-Speciﬁc Optimizations.
Researchers have found creative ways to implement individual operators or groups of operators in a more memory-efﬁcient 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-efﬁcient 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 (Bul ` o 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 signiﬁcant memory savings.
Although these hand-crafted techniques independently result in memory savings, there is no one-size-ﬁts-all recipe, and different implementations perform best on different architectures.
In contrast, MON E T automatically ﬁnds 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 ﬁnding 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 difﬁculties 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), MON E T is the ﬁrst system to jointly optimize a given architecture using both local and global techniques.
[Figure]
1 S N = {} /* Saved tensors for backward */ 2 L = { inputs, θ } /
* Local tensors for forward */ 3 for i = 1 ...N dox i = forward i ( L )
Add x i to L
Remove all tensors from L that are not used later if s Ni then
Add x i to S N 9 return L
[Figure]
P RELIMINARIES
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 forward i , and edges ( i,
j ) ∈
E specify the data-ﬂow 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 N j
= { i : ( i, j ) ∈ E } be the set of all incoming edges of an operation j .
We will ﬁrst 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
D k .
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 forward i ( · ) depends on a set of tensors L stored in local memory.
These tensors include model parameters Θ , computational dependencies N i , 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 forward i 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 x i are used in the backward pass, and have to be saved for later.
We use a checkpointing schedule s N to determine which.
Formally, s Ni ∈ { 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 forward i operator (Alg.
1 L. 4) allocates memory.
All other operators perform mere bookkeeping on existing tensor.
It is thus sufﬁcient to study the peak memory consumption m Ni in forward
i for each node i .
Let L i , S Ni be the set of local tensors L and saved tensors S while calling forward i respectively.
L i includes all parameters and computational dependencies for this and later forward passes L i =
Θ ∪ { x j : j ∈ N t for any t ≥ i and j < i } .
L i is constant and computed ahead of time.
The schedule s N determines the set of saved tensors S N i = { x j : s N j = 1 for j < i } .
In addition, each forward operator uses a certain amount of workspace memory c i to store intermediate results.
The total memory consumption of a forward operator is thus
[Equation]
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.
backward k ( · ) of each node k depends on a set of gradient tensors
ˆ L and forward tensors { x i : i ∈
D k } .
Any gradients required by the current and later backward passes are stored in local memory ˆ L .
Dependencies D k may either be stored in S k or need to be recomputed from checkpoints in S k .
Recomputation involves forward computation of one or more nodes, which increases computational overhead, and allows for a new set of tensors S k − 1 to be saved.
After recomputation, all dependencies D k 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 forward
i (Alg. 2 L. 7) and backward k (Alg. 2 L. 12) operations.
For the backward k operation, let ˆ c k be the workspace memory, ˆ L k be the set of gradient tensors stored,
D k
= { x i : i ∈
D k } be the forward tensors used, and S k − 1 be the set of newly saved tensors.
Here ˆ L k and D k can be pre-computed.
The total memory consumption for the backward k call is
[Equation]
[Equation]
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 forward i call (Alg. 2 L. 7) depends on the set of local tensors L , checkpoint dependencies
D , saved tensors S , and gradient tensors ˆ L , named L ki , D k , S k − 1 i , ˆ L k respectively.
Following the forward pass:
[Equation]
Unlike the forward pass, L ki is no longer constant, but depends on past saved tensors and future recomputations in ( s, r ) :
L ki =
Θ ∪ { x j : j ∈ N t for any t ≥ i with r kt = 1 and j < i } .
In the next section, we show how to take this formalization of the forward and backward pass, and ﬁnd an optimal execution plan including checkpointing schedule ( s, r ) , forward i implementations, and backward k implementations, under a ﬁxed memory budget.
M ETHOD
Our goal is to ﬁnd a global checkpointing schedule ( s, r ) and local forward
i /backward k implementations that jointly minimize the computation cost τ within a memory budget M .
We show how to express this optimization in a 0-1 integer program and efﬁciently solve it.
To this end, we linearize any peak memory consumption constraints, ensure that the checkpointing schedule is valid, and solve to minimize a computation cost objective.
We keep track of the three contributors to memory and computational cost - forward pass, backward pass, and recomputation of forward operators.
Memory Constraints.
Consider the case of basic checkpointing using only a single implementation for forward i and backward k .
The memory consumption of the forward 1 and backward 2 pass are linear in s , and thus efﬁciently expressed in an integer program.
However, recomputation depends both on s k − 1 and r k in a non-linear manner through the local memory L ki .
This joint dependence on optimization variables gives rise to quadratic constraints, which cannot directly be incorporated into an integer program.
For simplicity in this derivation, we bound the set of local tensors from above, assuming every future tensor is recomputed.
We give more information about this in Appendix B.
The upper bound ¯ L ki is constant, yielding a linear upper bound ¯ m ki of the recomputation memory ˜ m
ki analogous to Eq.
The set of memory constraints is thus m i ≤ M ∀ i and ˆ m k ≤ M ∀ k and ¯ m ki ≤ M ∀
k,i (4)
To enable operator optimization, we use a bit-vector δ to indicate the selection of an operator implementation.
We add δ to the constraints which allows us to jointly optimize checkpointing ( s, r ) and operator implementations δ .
Forward Operator Optimization.
Let each forward operator forward i have multiple different implementations
I i = { a, b, c, . . .
For examples, convolution may be implemented using matrix multiplication, the Winograd algorithm (Winograd, 1980), a Fourier transform, etc.
(Chetlur et al., 2014).
All implementations follow the same DAG structure, and thus use the same dependencies N i .
However, each implementation trades workspace memory { c ai , c bi , . . .
} for computational efﬁciency { τ ai , τ bi , . . .
} in a different manner.
Our experiments show that this trade-off is often complex.
Our goal is to represent the peak memory when using multiple forward i implementations in the forward pass and recomputation.
Let δ i,a ∈ { 0 , 1 } indicate that implementation a ∈ I
i is used for forward i in the forward pass.
Each forward operator should use exactly one implementation (cid:80) l
δ i,l = 1 .
The choice of implementation determines the operator’s computational cost (cid:80) l τ li δ i,l and workspace memory c
i = (cid:80) l c li δ i,l .
Analogously, each recomputation of forward i during backward k chooses between implementations δ ki,a ∈ { 0 , 1 } when needed (cid:80) l δ ki,l = r ki , with equivalent cost estimates (cid:80) l τ li δ ki,l and workspace memory use
c ki = (cid:80) l c li δ ki,l .
In this formulation, all additional memory requirements remain linear and are directly integrated into the linear memory constraints or their linear relaxations (equation 4).
Backward Operator Optimization.
Let each backward operator backward k have a set of different implementations
ˆ I k
= { a, b, c, . . . } .
Each implementation again trades workspace memory { ˆ c ak , ˆ c bk , . . .
} for computational cost { ˆ τ ak , ˆ τ bk , . . . } .
While gradient tensors follow the ﬁxed DAG structure, different implementations may depend on different forward activations {D ak , D bk , . . . } .
For example, in-place activated operators (Bul ` o et al., 2018) depend on their output activation, while regular operators use the input activation.
This change in the dependency structure makes optimizing for backward-operator implementations challenging.
We again aim to represent memory in terms of implementations for each backward k operator.
Let ˆ δ k,a ∈ { 0 , 1 } indicate that implementation a ∈ ˆ I k is used at node k in the backward pass.
Each backward operator should use exactly one implementation (cid:80) l ˆ δ k,l = 1 , with a computational cost (cid:80) l ˆ τ lk ˆ δ k,l and workspace memory ˆ c k = (cid:80) l ˆ c lk ˆ δ k,l .
The workspace memory adds a linear constraint to the memory consumption ˆ m k equation 2.
The biggest changes to the optimization problem, comes from the changing dependency structure .
D k is no longer constant.
Instead, the implementation of a backward operator changes the set of computational dependencies D k obtained from D lk .
To deal with this changing dependency structure, we use the indicator vector ˆ δ k to select memory contribution of dependencies from the chosen implementation.
This changes the backward memory consumption to ˆ m
[Equation]
= (cid:88)
[Equation]
l
[Equation]
and the corresponding peak recomputation memory ¯ m ki to
[Equation]
Note, the last term of equation 5 and equation 6 are quadratic in the original optimization variables s k − 1 i , which determines S k − 1 , and ˆ δ k,l .
However, for binary variables, it can be linearized using an auxiliary variable (see Appendix C.4).
We show the full equation expansion in Appendix C.1.
Checkpointing Constraints.
The computational dependencies of forward and backward operators impose strict constraints on the checkpointing schedule.
Any schedule violating these constraints cannot be executed.
Recomputation r ki requires saved s k − 1 j or recomputed r kj dependencies j ∈ N i , and only previously stored or recomputed tensors can be saved:
[Equation]
and s
[Equation]
≤ s k
[Equation]
Furthermore, all forward tensors D lk required by backward k need to be stored or computed
[Equation]
Objective.
Our goal is to minimize the amount of computation required for the forward and backward pass.
This is represented as the sum of computational costs of all operators:
[Equation]
Objective equation 9 with constraints equation 4, equation 7, equation 8, and deﬁnitions equation 1, equation 5, equation 6 form our ﬁnal optimization objective.
It jointly solves for the optimal implementation of each forward and backward operator, as well as an efﬁcient checkpointing schedule.
E XPERIMENTS
Implementation Details.
We develop MON E T in PyTorch v1.5.1 and solve the joint optimization problem using the Gurobi (2014) solver.
Appendix D provides more implementation details and a full list of optimized operators.
The UNet experiments use 608 × 416 inputs following prior work (Jain et al., 2019).
All other experiments use 224 × 224 inputs following conventions (Krizhevsky et al., 2012; Simonyan & Zisserman, 2015; He et al., 2016).
Batch size for the experiments is ﬁxed to be the maximum at which the model can be trained using baseline PyTorch on a 16 GB GPU.
Since Checkmate’s (Jain et al., 2019) execution engine is built for TensorFlow, and an ofﬁcial Gist (Jain et al., 2018) implementation is not available, we reimplement them in PyTorch for our comparisons.
Our Checkmate implementation is competitive, it uses the original Checkmate solver and has the same network structure as MON E T .
Checkmate does not optimize for operator implementations like convolutions, so we show its runtime using the default convolution algorithm (Checkmate-D).
For a stronger comparison, we also show the runtime of a Checkmate schedule that is post-optimized to greedily run the fastest convolution algorithm (Checkmate-O).
Wherever not explicitly speciﬁed, we compare with Checkmate-O.
All checkpointing schedules are run using the same software implementations and costs are proﬁled on the same hardware (NVIDIA P100 GPUs).
In order to compare against operator-speciﬁc optimizations, we reimplement all Gist techniques in PyTorch and run them on our execution engine.
See Appendix E for more details about our baseline implementations.
Detailed Comparison to Baselines.
(a) Checkpointing : Table 1 compares the memory savings obtained by MON E T and Checkmate for ﬁve different models when computational overhead over PyTorch is ﬁxed to be 10%.
MON E T schedules use 2-3 × less memory than PyTorch.
For the same computational overhead, MON E T uses 1.2-1.8 × less memory than Checkmate.
Fig.
3 shows more detailed runtime-memory trade-offs of MON E T to PyTorch and Checkmate for different models.
We plot the average iteration time of training as % overhead [Table] [Figure] [Figure]
GB [Figure] GB
[Figure] GB 8 GB 9 GB 10 GB ResNetCheckmate - 8. MONeT-NoOp 1.
MONeT 7. GoogleNet
Checkmate12.
MONeT-NoOp 0.
MONeT 3.
VGGCheckmate - - - 0. MONeT-NoOp - - - 0.
[Table]
MON E T and Checkmate schedules.
The memory budgets range from 5 GB to 10 GB, or equivalently, 0.33 × to 0.70 × PyTorch memory consumption.
Batch size for these models is mentioned in paranthesis.
For all models, MON E T reduces memory usage by 3 × (0.33 memory ratio) as compared to baseline PyTorch with 9 − 16% compute overhead.
For the same memory budget, MON E T schedules are up-to 34% faster than Checkmate schedules.
Note that we measure the empirical performance of the schedules running on GPUs instead of just providing a simulation of runtime and memory using the solver values; this is important since Checkmate does not consider workspace cost and overestimates its savings.
For networks with individual memory-intensive layers, like VGG-16, operator optimization becomes even more important for reducing memory; Checkmate can reduce memory for VGGonly up to 7 GB, whereas MON E T with its optimizations is able to run VGG-16 with 5.5 GB memory.
The small runtime improvement of MON E T schedules over PyTorch for VGG-16 and UNet at higher memory budgets is mainly because of choosing faster convolution algorithms.
MobileNetV2 uses depthwise convolutions, and hence does not signiﬁcantly beneﬁt from joint convolutionoptimization.
As a result, the performance of MON E T and Checkmate is closer for MobileNet-V2.
We provide additional results for MON E T on a memory-intensive model, 3D-UNet (C¸ic¸ek et al., 2016), in Appendix J, for which we observe a consistent memory reduction to 0.54 × of PyTorch memory with an overhead of 8.86%.
For our evaluation, we cap the solver time to 24 hours for both MON E T and Checkmate, and run the schedule thus obtained on our execution framework.
At tighter memory budgets for non-linear models like ResNet-50 and GoogleNet, Checkmate is unable to ﬁnd a feasible solution within a couple of hours.
In contrast to Checkmate, MON E T ﬁnds the execution plans efﬁciently.
For all the models and memory limits that we evaluate, MON E T reaches a 5% close-to-optimal solution within [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
a few hours or sometimes even minutes.
Table 2 shows the time it takes for the solver to reach 5% close to the optimal solution, for Checkmate, MON E T-N O O P
( MON E T with checkpointing enabled but operator-optimization disabled), and MON E T .
MON E T-N O O P converges to a close-to-optimal solution 1.6 × -117.4 ×
faster than Checkmate.
For larger models, MON E T ’s solver converges to a close-to-optimal solution up to 27 × faster than Checkmate.
Note that running a solver is a one-time cost for a model - once a MON E T schedule has been solved for, it can be used by everyone to train the model for different purposes with different batch sizes.
The cost (typically seconds to hours) is tiny compared to the efforts and costs to develop a model for distribution in most cases.
See
Appendix H for more discussion regarding solver times, problem statistics, and full Table 2 data.
(b) Operator optimizations : Table 3 shows the comparison of MON E T with Gist. While MON E T can determine a range of memory-runtime tradeoffs, purely operator-optimization-based schemes like Gist only provide a single memory-runtime data point.
For MON E T , we show the memoryruntime data point with the most memory saving. MON E T uses 1.4 × -2.1 × less memory than Gist for multiple architectures while maintaining full-precision. Overall, Gist provides impressive memory savings, but incurs a high computation cost to achieve the savings.
While we get similar memory saving results for reimplemented-Gist as Jain et al. (2018) for VGG16, our compute overhead results are higher. This could be because of evaluations on different frameworks (PyTorch v/s CNTK) and different GPU models (Nvidia P100 v/s Nvidia Maxwell
GTX Titan X).
Gist uses dense to sparse conversion using cusparseSdense2csr in one of its techniques. For the ﬁrst ReLU-Conv layer in VGG-16 (shape (2207744,256) ), this function takes 144ms, which itself is 10% of the VGG-16 execution time. We see similar results for other networks. To ensure a fair comparison, we focus on the maximum memory savings obtained by
MON E T with Gist, while reporting the compute overhead for completeness.
Ablation Experiments.
Fig.
4 shows additional ablation experiments.
We show the % compute overhead over PyTorch on ResNet-50, GoogleNet, and VGG-16 for different types of MON E T checkpointing schedules with a memory budget of 8 GB - with no operator optimizations enabled, with only one type of operator optimization enabled (conv-optimized, output-activated optimized, intermediate-activated optimized), and with all optimizations enabled. Schedules which do not jointly optimize convolution algorithms are run with greedily post-optimized convolution algorithm. Plots for other models look similar to that of ResNet-50 and GoogleNet. The only difference between ’none’ and ’conv’ is that convolution algorithms are jointly optimized in the latter.
However, this fact leads to signiﬁcant improvement in compute time for all cases. In fact, convolution algorithms have complex workspace memory - compute characteristics, reserving slightly more memory for convolution workspace while checkpointing can allow for a much faster convolution (see Appendix I). This makes it important to jointly optimize conv algorithms with checkpointing. Similarly, output-activated optimization also provides signiﬁcant beneﬁts over vanilla checkpointing, since it effectively reduces the number of recomputations required.
For memory-intensive net-
PyTorch (14.7 GB)
MONeT ( 8.0 GB )
M e
[Figure]
( M
B )
PyTorch (860 ms)
MONeT-NoOp (939 ms)
MONeT (908 ms) −−  ·
[Figure]
i
[Figure]
[Figure]
[Figure]
Most happens works, intermediate-activated optimization becomes more important.
Jointly optimizing all strategies together gives the least computational overhead.
See Appendix G for detailed ablation plots.
Detailed Case Study.
The top graph of Fig.
5 shows memory usage while executing PyTorch,
MON E T without operator optimization, and MON E T for ResNet-50 at batch size 184.
As the training progresses along network layers represented on X-axis, PyTorch and both MON E T schedules store forward-pass outputs, leading to an increasing memory footprint.
MON E T reaches peak memory of 8 GB, whereas PyTorch requires 14.7 GB.
Stored forward outputs are freed up one after other as backward pass proceeds, leading to reduced usage of memory.
According to the checkpointing schedule, MON E T saves only a subset of the outputs stored by PyTorch, resulting in the memory saving shown in the middle graph for layer outputs that are not stored.
The bottom graph shows the per-layer compute overhead of recomputation of MON E T over PyTorch.
For MON E T , later layers which are backward operators result in a recomputation of the forward, and have higher overhead.
C ONCLUSION
We present MON E T , a system to automatically reduce memory requirements for training deep networks.
MON E T jointly optimizes local (operator-level) and global (graph-level) optimizations to yield a compute- and memory-efﬁcient checkpointing schedule.
MON E T reduces memory usage by 3 × over PyTorch, with a 9 − 16% compute overhead.
It uses 1.2-1.8 × less memory than the state-of-the-art automated checkpointing framework for the same computational cost.
Our experimental results show that MON E T leads to better memory-computation trade-offs compared to the state-of-the-art.
A CKNOWLEDGMENTS
We would like to thank the anonymous reviewers for their feedback.
Aashaka Shah and Vijay
Chidambaram were partially supported by donations from VMware and Google.
Chao-Yuan Wu was partially supported by a Facebook Fellowship.
Jayashree Mohan was supported by a Microsoft
Research Fellowship.
The results presented in this paper were obtained using the Chameleon testbed supported by the National Science Foundation
A N OTATIONS
Table 4 gives a list of notations used in the paper along with explanations.
Notation
Meaning
Tensors x i
[Table]
y k
[Table]
Sets S Ni [Table] L i
[Table]
D k
[Table]
S k −
[Table]
ˆ
L k
[Table]
S k [Table] L ki [Table]
I i Set of implementations for operator forward
[Table]
ˆ I k
Set of implementations for operator backward k . Solver variables s Ni Indicate if output of forward i is stored in memory in the forward pass.
{+s k − 
 i+}
Indicate if output of forward i is stored in memory when computing backward k .
r ki
Indicate if forward i is recomputed before computing backward k .
δ i,l
Indicate if forward i uses implementation l ∈ I i in the forward pass.
δ ki,l
Indicate if forward i uses implementation l ∈
I i when recomputed before backward k .
ˆ δ k,l
Indicate if backward k uses implementation l ∈ ˆ I k .
Memory formulations m i
Peak memory of forward i in forward pass.
¯ m ki
Peak memory of forward i when it is recomputed before backward k .
ˆ m k
Peak memory of backward k .
Operator costs c li
Workspace memory of operator forward i executed using implementation l ∈
I i .
ˆ c lk
Workspace memory of operator backward k executed
[Table]
B B OUNDS ON LOCAL MEMORY
In Section 3, we mentioned that local memory L ki is dependent on solver variable r kt .
L [Equation]
[Equation]
j : j ∈ N t for any t ≥ i with r kt = 1 and j < i } .
In order to remove this dependence, we can get an upper bound ¯ L ki on L ki by assuming that all future tensors after i will always be recomputed, that is, r kt = 1 , ∀ t > i , and
[Equation]
Our experiments also use this upper bound.
It is possible to tighten the upper bound by noting that r kt may be 1 only in the case when t ≤ k .
That is, forward node t will not be recomputed before computing backward of node k
if node t lies after node k .
Thus, a tighter bound to L ki follows
L k i ⊆ ˙
L k i =
Θ ∪ { x j : j ∈ N t for any t ≥ i and t ≤ k and j < i } ⊆
¯ L k i .
C D ETAILED CONSTRAINTS C. E XPANDED BACKWARD PASS MEMORY CONSTRAINTS
Sec.
4 formulates backward peak memory ˆ m k and recomputation peak memory ¯ m ki as sum of memory of a set of tensors.
We expand the memory formulation and represent it in the terms of optimization varaible here:
[Equation]
C. C OMPLETE MEMORY CONSTRAINTS
In this section
[Equation]
we present the complete memory constraints which we use for MON E T optimization.
These constraints include the recomputation variable r ki , which was excluded from the main text to make understanding simpler.
As discussed in Sec. 3, the peak memory of a forward i recomputation before computing backward k is denoted by ˜ m ki .
This represents the recomputation memory (renamed to m kRi ) when forward i is actually recomputed, that is, r ki = 1 .
When this is not true, the peak memory ( ˜ m kSi ) only depends on stored checkpoints S k − 1 i , checkpoint dependencies for D k , and gradient tensors ˆ L k .
Thus,
[Equation]
Local memory L k can be bounded by ¯ L k , which gives us ¯ m kRi .
To add forward operator optimizations to ¯ m kRi , we recall the trade-off between workspace memory and compute time.
We replace the workspace memory contributor r ki c i in equation 12 with (cid:80) l δ ki,l c li .
The complete memory constraints are: m i ≤ M ∀
i and ˆ m k ≤ M ∀ k and ¯ m kRi ≤ M ∀ k,i and ˜
[Equation]
C. I N - PLACE CONSTRAINTS
We show how to represent the decision of computing an operator using an in-place or out-of-place implementation.
If an operator like ReLU uses an in-place implementation, its input tensor is overwritten with its output.
In this case, its input tensor cannot be stored or used as input to a computation in this stage.
This needs to be reﬂected in our constraints.
We introduce two new binary variables to model in-place computations:
q k i represents if forward i is recomputed in-place when computing backward k .
p ki represents that the output of forward i has been computed and will not be overwritten by any other forward node recomputations in this stage.
If q ki is true, then p kj will be false else p kj will be the same as r kj , where j ∈ N i .
Further, s k − 1 j will also be false if q ki is true.
This can be written in the form of boolean constraints as follows:
[Equation]
and
[Equation]
and
[Equation]
The checkpointing constraint 7 changes, with p kj replacing r kj on the RHS.
Further, q ki (or p kj ) can only be true if forward i (or forward j ) is actually recomputed prior to computing backward node k.
Thus, p kj ≤ r kj and q ki ≤ r ki .
C. C ONSTRAINT L INEARIZATION
The memory constraints we introduce in Section 4 contain quadratic terms in the form of x i · x j , with x i , x j ∈ { 0 , 1 } .
The quadratic terms cannot directly be incorporated into an integer program.
However, we can linearize these terms by replacing each quadratic term x i · x j by an auxiliary variable α i,j ∈ { 0 , 1 } and introducing additional linear constraints α i,j ≥ x i + x j − 1 , α i,j ≤ x i , and α i,j ≤ x j .
After this substitution for all quadratic terms, all constraints in MON E T are linear.
D I MPLEMENTATION
We develop MON E T in the PyTorch (v1. 1) framework.
We use PyTorch’s default Autograd package for backward implementation of elementary functions when the autograd implementation is stateless.
In all other cases, we implement custom forward and backward functions leveraging PyTorch ATen library functions to ﬂexibly support multiple operators and execution schedules.
Each backward operator implementation is annotated with its computational dependencies, which is generally the input or the output of its corresponding forward operator.
Certain backward operators implementations may have dependencies on intermediate activations generated in the forward pass.
For example, an intermediate-activated ReLU backward uses an encoded bitmask representing the sign of forward operator’s input.
We annotate this as an intermediate storage node and add it to our optimization problem, with a strict recomputation dependency of the interemediate storage node on its creator node.
Our operator optimizations select from different backward operator implementations, convolution algorithms, in-place operators etc.
We split the convolution backward operator into two - a parameter-gradient operator followed by an input-gradient operator.
Since the inputgradient operator does not have any computational dependency on the forward pass, we can agressively free the forward input tensor right after the parameter-gradient is computed.
We also reuse
BatchNorm statistics in case of their recomputation.
For our experiments, we limit ourselves to full precision training as quantization or lower precision computations introduce additional noise into
SGD and change its convergence properties.
We solve the joint optimization problem using the
CVXPY (Diamond & Boyd, 2016; Agrawal et al., 2018) solver with Gurobi (2014) backend.
MON E T workﬂow.
We obtain the forward pass dependencies in MON E T by JIT tracing a model to obtain its graph.
We proﬁle each layer for workspace memory and compute cost, and obtain memory usage of the tensors from their shape and type.
Note that the workspace memory for many convolution operators in VGG-16 is greater than 2GB, making it an important factor to model.
Unlike prior approaches like Checkmate, we account for this workspace memory in our optimization problem, bringing the memory model very close to actual memory allocation.
We phrase a boolean integer programming problem using the generated graph and the proﬁled compute cost and workspace memory and solve it using the CVXPY (Diamond & Boyd, 2016;
Agrawal et al., 2018)
modeling language and GUROBI (Gurobi, 2014) solver.
The solution is used to generate a schedule that can be run by the MON E T scheduler.
Operator optimizations.
We divide operator optimizations according to the different type of implementations they select from.
(1) Output-activated : Backward calculation of operators like ReLU and
BatchNorm can have computational dependency either on on their forward node’s inputs or outputs.
(2) Intermediate-activated : Backward of ReLU has computational dependency on a 1-bit encoding of the sign of its forward node’s input.
Backward of MaxPool is calculated using an intermediate 8bit output-shaped tensor which contains the kernel-index of the maximum element.
(3) Convolution algorithms : We choose from 8 forward and 6 backward cuDNN convolution algorithms.
(4) Inplace operations : The solver can choose to do inplace computation for operators like ReLU forward.
We discuss constraints for in-place operator selection in C.3.
All MON E T experiments enable in-place operation selection.
E B ASELINE IMPLEMENTATIONS
Checkmate implementation.
We reimplement Checkmate (Jain et al., 2019) in PyTorch for our comparisons.
We use the Checkmate solver as-is to obtain Checkmate schedules.
Since Checkmate does not provide an execution engine for PyTorch, we run the generated Checkmate schedules on our own execution framework.
Our inference engine uses the same operator implementations for
Checkmate and MONeT. We have released our Checkmate implementation with the MON E T code.
Gist implementation.
Gist (Jain et al., 2018) is an operator-based memory-efﬁcient scheme for training DNNs.
It encodes stashed forward tensors into smaller tensors which require less memory. Jain et al.
(2018) evaluate Gist using CNTK on an Nvidia Maxwell GTX Titan X GPU.
Since we implement MON E T in PyTorch and have access to an Nvidia P100 GPU, a direct comparison with the numbers in the Gist paper is not possible. As an ofﬁcial Gist implementation is not available, we reimplement it on PyTorch and evaluate its execution using MON E T ’s execution framework.
We implement all Gist optimizations — Binarize (intermediate encodings for ReLU-Pool layers),
Sparse Storage Dense Compute (compress and store sparse convolution inputs in ReLU-Conv layers as sparse storage), Delayed Precision Reduction (storing stashed non-compressed tensors in FP-16, but computing in FP-32), and Inplace (performing ReLU operators in-place wherever possible) over
MON E T ’s execution framework.
In Gist, the Sparse Storage Dense Compute (SSDC) technique creates a sparse storage tensor in the Compressed Sparse Row (CSR) representation using the Nvidia cuSPARSE library.
The dense storage is reshaped into a 256-sized column tensor before storing it in a sparse format, allowing the column index of CSR representation to be saved in 8 bits instead of using 32 bits (termed Narrow Value Optimization in the paper).
We also implement SSDC using
Nvidia’s cuSPARSE library (function cusparseSdense2csr ) with CUDA Toolkit version 10.
using PyTorch’s C++ extensions.
In their paper, Jain et al.
(2018) use the most memory-efﬁcient convolution algorithms in Gist and compare its memory saving against a baseline which also chooses the most memory-efﬁcient convolution algorithm. Using memory-efﬁcient convolution algorithms, our Gist reimplementation can train VGG-16 with 0.55 × of the PyTorch-required memory (1.81 × memory footprint), which is close to the data presented by Jain et al.
However, it is 59% slower than when convolution selection is enabled, in which case it can train using 0.76 × of the PyTorch-required memory.
Since implementing Gist using memory-efﬁcient convolutions is not optimal in terms of compute time, we implement Gist to use PyTorch’s convolution selection algorithm.
For all models other than VGG16 and UNet, we see similar memory savings for Gist with memory-efﬁcient convolutions and with convolution-selection enabled. We have released our Gist implementation with the MON E T code.
F O N OPERATOR SELECTION FOR CHECKMATE
In this section, we brieﬂy explain the difﬁculties of including operator selection directly into checkmate (Jain et al., 2019).
We will refer directly to notation and equations in the checkmate paper (arxiv v3; 14 May 2020).
The most direct way to incorporate operator selection into checkmate is to introduce an auxiliary variable R vt,i ∈ { 0 , 1 } that refers to re-computing layer i at time t using implementation v .
Most constraints in equation 1 could stay the same, given R t,i = (cid:80) v R vt,i , and loss (1a) (cid:80) t (cid:80) i (cid:80) v R vt,i C vi .
Some of our operators produce a different kind of checkpoint (e.g. binary activated ReLUs), which could be handled in check-mate by splitting S vt,i .
The main issues in checkmate arise in the memory modeling and its relaxations (eq 4,5,7).
The memory consumed by a speciﬁc checkpoint may depend on the operator implementation: DEPS[k] and
USERS[i] both depend on the operator implementation (output activated, input activated, ...).
In short, the checkmate computation graph is dynamic and depends on operator implementations.
The most direct way to address this is to mem freed t ( v k ) =
(cid:80) v R vt,i mem freed t ( v k ) in a implementation dependent way mem freed v t ( v k ) , and select the right version dependent on the operator used.
Likewise, we need to extend FREE vi,t,k to account for different operator implementations in R vt,k .
Likewise the product in equation (5) will now go over all implementations R vi,j using different
USERS sets.
This leads to a linear blowup in the number of constraints, and number of auxiliary variables, leading to an at least quadratic expansion on computational costs.
Furthermore, mem freed t ( v k ) =
(cid:80) v R vt,i mem freed t ( v k ) is a quadratic constrain that further needs to be resolved using additional auxiliary variables.
Given that Checkmate already pushes the limits of [Figure] [Figure]
ov
[Figure]
Int
[Figure]
current solvers, it is unlikely able to handle this explosion in constraints and variables, without signiﬁcant modiﬁcations. MON E T in the other hand represents the compute-graph more compactly and efﬁciently integrates different operator implementations.
G D ETAILED ABLATIONS
Fig. 6 shows a detailed plot of our ablation experiments comparing the compute overhead of variants of MON E T across a range of memory limits. Y-axis shows the compute overhead over PyTorch and X-axis shows the memory ratio to a PyTorch model.
All variants which are not conv-optimized are greedily post-optimized to use the fastest convolution. We see that MON E T with no operator optimization (NoOp) is generally slower than the other variants for all models and memory limits. Convolution and output-activated optimizations are both important in reducing compute overhead. MobileNet-V2 uses depthwise separable convolutions, and hence does not signiﬁcantly beneﬁt from convolution-optimization.
Further, MobileNet-V2 has hardtanh operators instead of ReLU operators, for which we have not implemented intermediate-activated backward optimization. Interemediate-activated optimizations provide memory savings in memory-intensive models, allowing models like VGG-16 to reach memory savings which are not attainable by other optimizations.
All optimizations together result in the least compute overhead for any model or memory limit.
H S OLVER TIME AND ILP STATISTICS
Solver runtime.
MON E T ’s solver runtimes vary for different models and different memory limits.
We evaluate schedules obtained using solver times set to a maximum of 24 hours.
For moderate memory limits, both MONeT and Checkmate achieve an optimal solution before 24 hours.
For tighter memory limits, the solution obtained by MONeT and Checkmate may not be most optimal.
For multiple models and memory limits, Table 2 in Sec. 5 shows the time it takes for the solver to reach 5% close to the optimal solution for Checkmate, MON E T-N O O P ( MON E T with checkpointing enabled but operator-optimization disabled), and MON E T .
We add the data for MobileNet-V2 and
UNet in Table 5 which also follow a similar pattern.
We also provide Table 6 which shows the time taken by the solver to reach 2% close to the optimal solution.
We note that it has similar behavior as the time taken the solver to reach 5% close to the optimal solution.
MON E T-N O O P converges to 2% close-to-optimal solution 1.3 × -139 × faster than Checkmate.
For larger models, MONeT’s solver converges to a 2% close-to-optimal solution up to 16 × faster than Checkmate.
At tighter memory limits for MobileNet-V2, the Checkmate solver reaches 2% close-to-optimal solution faster than
MON E T , but is still much slower than MON E T-N O O P .
ILP statistics.
For different models, Table 7 shows the solver statistics after presolving for the problem formulated by Checkmate, MON E T-N O O P and MON E T for a 10 GB memory limit.
It shows the number of forward operators in the model and the number of constraints and variables for each solver.
MON E T-N O O P , which is MON E T with only checkpointing enabled and without using operator optimization, has on average 50% fewer constraints and 67% fewer variables than
Checkmate
Jointly-optimized MON E T has a slightly larger number of constraints, and on average [Table] solution 1.6 × -117 × faster than Checkmate. MONeT gets close to 5% of the optimal solution only in a few hours, and up-to 16 × faster than Checkmate for larger models.
[Table]
to reach 2% close to optimal solution.
MONeT-NoOp reaches a 2% close-to-optimal solution 1.3 × -139 × faster than Checkmate. MONeT reaches a 2% close-tooptimal solution within few hours in most cases, and up to 27 × faster than Checkmate for larger models. 40% fewer variables than Checkmate. MON E T ’s formulation is more efﬁcient, and might be the reason that it reaches a good solution faster than Checkmate.
[Table]
I C ONVOLUTION A LGORITHMS
Fig.
7 shows the complex workspace memory-compute trade-off for different convolution algorithms.
The memory used is not always inversely proportional to the compute requirement.
Jointly optimizing convolution algorithms enables MON E T to make the best decisions for which convolution algorithm to select.
J A PPLICABILITY TO MEMORY - INTENSIVE MODELS
To further show MON E T ’s applicability to memory-intensive models, we evaluate it on 3DUNet (C¸ic¸ek et al., 2016), a fully-convolutional model for volumetric images.
Fig.
8 presents the runtime-memory trade-off for MON E T on 3D-UNet.
We used a commonly used 3DUNet implementation (Wolny, 2019; Wolny et al., 2020) with training conﬁguration similar to 3DUnet confocal boundary provided in the repository and a batch size of 22, which just ﬁts on a 16 GB P100 GPU.
MON E T reduces memory usage to 0.54 × of PyTorch, while incurring 8.86% overhead in compute time.
At a memory ratio of 0.81, MON E T incurs almost no computational overhead, because it makes use of operator optimizations and is able to bring down the recomputation cost to zero.
[Figure]
[Figure]