Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: Optimizing Neural Network Memory is broadly done through two channels - 1)  local operator level - like storing only signs of a ReLU function or bit quantization 
2) global graph level - like optimizing checkpointing schedule for a given compute graph. These channels are usually orthogonal/independent. 

In this paper, the authors propose MONet which tries to find the best checkpointing schedule that can jointly optimize both the above channels. 
The authors create an auxiliary graph to encapsulate operators and perform schedule optimization on the new graph rather than the usual graph in existing frameworks.

To obtain the new graph, they first perform a theoretical quantification of the peak memory consumption of the forward pass,backward pass and recomputation.
Then, under a fixed memory budget M, we try to optimize the operators for computational efficiency.

Each operator has multiple implementations that trade-off workspace memory 'c' and compute efficiency \tau and we can exactly use only one implementation per operator. 
So we end up with a linear combination of compute costs for all operations that looks like 

\sum_i\sum_l(\tau_i^l \delta_il) where \delta is an indicator of which implementation to choose.

With 3 such double summations, the authors come up with a join objective function which constraints the overall memory sum to <=M 
and minimizes the compute cost objective above (which can be solved by a linear program).

On several popular architectures like ResNet-50, VGG-16 and UNet etc, MONeT outperforms Checkmate by abt 1.3-2x on memory and full-blown PyTorch by abt 2-3x.

Pros:
1. Thoretically optimized
2. Experiments on multiple CNN architectures

Cons:

1. Effetcively no comparisons against any other schemes except Checkmate.
2. No other architectures optimizations like Transformers or large scale FFNs.
review_title: The paper proposes MONeT to jointly optimizing both local operators and the global compute graph. The authors express this joint optimization problem as a integer program. Results are impressive with 3x memory reduction at a small overhead or 1.2-1.8x reduction at the same computation.

final_score: rating: 8: Top 50% of accepted papers, clear accept
review_content: 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.

The memory constraints are carefully analyzed for the forward and backward passes, and expressed as a 0-1 integer program, which is then solved using the state-of-the-art solver Gurobi. The experimental evaluation confirms the theoretical gains provided by the solution of the optimization problem. In particular, the authors compare with a vanilla implementation in PyTorch, with the Checkmate-D (default Checkmate that uses global techniques), and with Checkmate-O (post-optimized to greedily run the fastest convolution algorithm). It is interesting to notice that MONET offers significant gains in all cases, even over Checkmate-O, underlying the need to perform the joint optimization in order to obtain the best schedule.

I advocate for the acceptance of the paper. The memory bottleneck is an acute problem in deep learning, and the paper provides a practical solution that can alleviate this problem to some extent. I encourage this line of work and hope to see such schedule optimization become a standard option in deep learning frameworks. The paper is very well written, with a clear description of the approach and convincing experimental evaluation, while providing abundant references about related work.

Minor comments, questions for the authors:
The memory requirements for some of the networks shown in the experimental evaluation are still small. While I understand that the GPU memory is limited to 16GB and there was a desire to compare against vanilla implementations, I think it would be interesting to show how models that require a lot of memory for training can benefit from MONET. For example, the UNet that was used in the paper seems to be the 2D version. The 3D UNet can require up to several hundred GB without optimizations (for varying batch sizes). Would it be possible to include even a simple estimate of the possible gains when using MONET? What would be the computational tradeoff to train it on the GPU used in the paper?

review_title: Valuable contribution for deep learning training

final_score: rating: 7: Good paper, accept
review_content: The authors present MONeT, an automatic approach to jointly optimize operator cost and checkpoint scheduling for deep learning on a fixed memory budget.  The paper thoroughly defines the problem, relevant previous work, and the MONeT framework.  Given a fixed GPU memory budget, MONeT solves an integer program in order to jointly minimize the computational overhead of checkpointing with various operator implementation.  This approach is intuitive, as previous approaches, such as the recently proposed CheckMate, only optimize the checkpoint schedule.  The derived integer program is also a nontrivial extension of previous work.  With MONeT implemented in PyTorch, a large number of empirical results are presented, which show the superiority of MONeT compared to CheckMate, and show memory savings (versus impressively slight overhead) compared to PyTorch.

The paper is well written, the description of computational cost and the derivation of the integer program were interesting, and results are very compelling (and easy to understand).  One area where the paper could be improved is on the notation used throughout Section 4 (see below for suggestions), which was difficult to follow due to the density of the section, the large number of different variables/variants of variables described, and some implicit definitions.  There are also a few small details and discussions which seem warranted, but, overall, I enjoyed reading this paper.

Specific comments:

-"n Checkmate, changes in operator implementation induce a different computation graph" <- While this is
technically true, only the cost associated with that operator changes,
yes?  In this way, Checkmate could run multiple passes over the
computed static graph with different operator costs, but this approach
would require an expoential (in the number of operators with varying
costs) number of evaluations.

-"We reimplement Checkmate in PyTorch" <-This is non-trivial, so please include some re-implementation details.

-Please include the version of PyTorch which was forked for the Monet  and Checkmate implementations in Section 5.

-As CheckMate currently only supports tensorflow, it would be very helpful if the authors could also release the source for CheckMate in PyTorch when the source for MONeT is released.

-Could you comment on the solver runtime needed to solve the integer
program in monet, in contrast to the solver runtime for the MILP in
checkmate?

-How do open-source solvers compare to the runtime for the commerical
Gurobi solver for the integer programs solved in monet?

-"We measure the empirical performance of the checkpointed schedules
running on GPUs instead of just providing the solver values; this is
important since Checkmate doesn’t consider workspace cost and
overestimates its savings... Hence,
we show the results with solver running for 1 day for both MON E T and
Checkmate. In contrast, MON E T finds the execution plans efficiently,
its 1-hr solution already close to the 1-day solution with a small
difference of 1-2%."  <- This paragraph of worded somewhat vaguely; is
this to say 24 hours were included in execution time?

-For VGG-16 in the ablation study in Figure 4, PyTorch exceeds
device-memory for this dataset, yes? Is this the reason why monet
achieves negative overhead; i.e., faster execution time than PyTorch
itself?

Comments/questions regarding notation:

-What is variable r in the schedule (s,r)?  It *seems* like r is the
 indicator function for activations which require recomputation.  Is
 that correct?  If so, please (please) state this explicitly in the
 paper.  If not, please define r.

-A supplementary table for variables would be helpful.  It took some
 time to find a definition for $y$ in Equation 2.  While I eventually
 found one in Algorithm 2 (please define this explicitly inline,
 preceding Equation 2), a table would have made this much easier,
 especially considering how dense Section 3 is.

-By the time the reader gets to Section 4, the table of variables
 becomes mandatory (there are different values of L and S with various
 superscripts, subscripts, and hats, it is very difficult to recall
 which is which and to look back in the dense text for their
 definitions).

-Also, if possible, please standardize notation by the three
 categories: 
(a) peak
 forward pass memory consumption
(b) peak backward pass memory consumption
(c) peak recomputation memory consumption
so that a reader can ascertain what collection Ls, Ds, and Ss are
being referred to.  I understand there is overlap between these three
categories, but there must be some organizational way to more
easily refer to these variables without having to research for their
definitions when reading later sections of the paper.


review_title: Intuitive approach and framework to push the state-of-the-art in memory-constrained deep learning

final_score: rating: 7: Good paper, accept
review_content: ## Summary:

The paper provides a framework (MoNET) to perform automatic memory optimization targeting deep neural networks. Their technique jointly optimizes the checkpointing schedule and the choice of individual ops to find the implementation with the least possible runtime overhead under certain memory constraints. They formulate the problem as a ILP problem where the objective is to minimize runtime subject to strict memory constraints. They show that the solutions achieve 3x reduction in memory footprint compared to a pytorch implementation with minimal runtime overhead.

## Strengths:
* Making joint decisions on checkpointing and on the implementation of individual ops. This enables the authors to make both global and local decisions and to exploit the synergies in both.
* Experiments on multiple NN models to show the impact of their technique.
* Easy to read formulation of the ILP problem.

## Weaknesses and Questions for authors:
* Solution times and scalability
The authors do not explicitly mention the solution times for their ILP formulation. Also, it is unclear to me how the ‘checkmate’ framework can be slower compared to MoNET (as authors claim in page 7), since checkmate is only solving for the optimal checkpointing schedule.  I am also curious to know the scalability of the system for larger problems, specially since the authors are doing a joint optimization. Statistics about ILP problem sizes would be helpful.
* Questions about the formulation
  * Why do you need to compute $r_i^k$ for $i > k$? The gradient computations of op $k$ only depends on ops with numbers > $k$ (Algorithm 2)
  * How do you get the computational efficiency of ops? If fusion of ops happens inside the compiler, then adding up the computational efficiency of individual ops may not realistically model the execution behavior.
  * The formulation does not take into account possible compiler optimizations such as fusion, tiling etc. for both memory consumption and runtime calculations.
* Questions about experiments
  * It’s unclear what implementations of ops are used in the PyTorch implementation? Are these default ops or did you do a post-optimization similar to what you did for checkmate?
  * What if you set memory ratio = 1 and compare it with the pytorch implementation, what is the overhead?

review_title: Good paper; please answer the questions posed 

Meta-review: The reviewers all agree that Monet proposed in the paper which optimizes for both local and global memory saving in Deep learning models is theoretically sound and experimentally convincing.
Accept!
Decision: Accept (Spotlight)