Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: Summary: This paper proposed a simple yet effective greedy algorithm with a new heuristics on checkpointing deep learning models so that people could train large model with restricted GPU memory budgets. The proposed method operates in an online setting and do not need static analysis of computation graph, thus could be used for both static and dynamic models. In a restricted model setting of linear forward network and equal space and time cost for each node, the author proves the proposed method could reach the same bound on tensor operation and memory budget with previous static checkpointing methods. The author also establish a theorem on tensor operation numbers between the proposed dynamical method and an optimal static checkpointing algorithm. In experiment, the author compared the proposed method with static techniques including the optimal Checkmate tool of Jain et al. (2020), showing the proposed method gives competitive performance without static model analysis in prior. The author also compared the proposed heuristics with prior arts on several static and dynamic models. Finally, the author described a prototype of PyTorch implementation of the proposed method. 

Pros: 
1. While goes under a limited setting, the theoretic analysis on the tensor operation and memory budget bound of the proposed method, as well as on the relationship between the proposed method and optimal static analysis method is novel and interesting. The experiment also shows the competitiveness of the proposed method by comparing to static methods. 
2. The author does a great job explaining the idea, concepts, procedures and experiments. 

Cons: 
1. The author provides the comparison between the proposed heuristics and others with the same greedy algorithm, but it seems not to have the full comparison to other dynamic checkpointing approach(e.g. Peng et al. (2020) ). Although experiments in the paper shows competitive results with static model, the main use cases of the proposed method might still be in dynamic models as in normal static model use cases, the time overhead of static analysis could be ignored compared to actual model training time.
2. As the proposed heuristics bears some similarity with the one used in Peng et al. (2020), it would be more convincing to also have an ablation study of replacing the heuristics used in Peng et al. (2020) with the proposed one. 

References:

[1] Jain, Paras, et al. "Checkmate: Breaking the memory wall with optimal tensor rematerialization." Proceedings of Machine Learning and Systems 2 (2020): 497-511.

[2] Peng, Xuan, et al. "Capuchin: Tensor-based GPU Memory Management for Deep Learning." Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems. 2020.

review_title: Interesting work

final_score: rating: 6: Marginally above acceptance threshold
review_content: The paper presents an online algorithm for dynamic tensor rematerialization.  Theoretically, it shows the same asymptotic order on the memory budget and tensor operations as of the optimal static approach.  By simulation, it shows the performance matches optimal static checkpointing in a few models.  A PyTorch prototype is implemented, which shows benefits of reducing memory footprint and increased batch size comparing with basic PyTorch models without checkpointing.

Merits of the paper:
- Address an important practical problem on how and when to perform checkpointing during DL training.
- Cover a pretty comprehensive study across theory, simulation and system implementation.
- The suggested system implementation looks simple and clean.
- Clearly written paper, which is easy to understand and follow.

Places to improve:
- The need of having dynamic approach in this area is not very well motivated.  Since the computation in most of DL models is repetitive over iterations, static approach would work pretty well.  Furthermore, since many models take long to train, spending minutes on analyzing static graph and obtaining an optimal solution seems to be time well spent, comparing with a suboptimal dynamic solution.  The motivation of developing dynamic approach needs to be strengthened.  

- Although some comparison with related work is done by simulation, really system evaluation is relatively weak - it only compares with the strawman PyTorch models without checkpointing.   It would be a lot more convincing with the results comparing with basic checkpointing approach (e.g., layer wise checkpointing) and some related work.

- For some of the reported models, like Unet, the approach reduces memory footprint and helps increase batch size, which however reduces  throughput.  These are probably not good examples to show the benefits of the approach.  But it might be beneficial to point out the underlying reasons for decreased throughput so readers know when the approach would/wouldn't work well and why. 
review_title: A useful and practical solution for dynamic checkpointing but system evaluation needs to be well strengthened

final_score: rating: 7: Good paper, accept
review_content: Contributions: a) analyzes multiple heuristics for which tensors to evict where compute overhead of rematerialization is minimal overall, b) suggested approach is just-in-time, and thus does not require any static analysis of the network. That is, unlike prior work in this area, it covers any network type with no prior knowledge, c) offers a good formal analysis of proposed heuristic in terms of its components: staleness, memory capacity and recursive replay cost - their formalization covers previously published heuristics as well.  Experimental framework is sound. And some encouraging results are shown delivering memory capacity saving of 30% to 90% with training slowdown of 2x or less.

Prior such work all required static analysis and planning of the network - and hence were of limited use. Significant contribution of this work is summed up at the end of Sec 4.3. It achieves 'as good' results as prior state-of-the-art (Checkmate, published at MLSys) - but without any prior knowledge of the model. This significantly increases the practical significance of this work. NeuIPS 2019 work of Kumar is the other often-cited work, also based on static planning, where the authors had already noted the primary limitation keeping it from getting adopted: "algorithm yields asymptotically better schedules, the schedule length and memory depend exponentially on the path width". 

I am also happy to see authors offer hardware detail of their experimental platform (Figs 2 and 4).  And PyTorch software prototype should make it easier to follow through in other frameworks. There is a decent variety in the chosen set of benchmarks as well.

review_title: This submission is based on the observation that intermediate activations can be replayed (through local forward prop calculations), rather than kept in memory for backprop to reach them after the full forward path calculation. As a result, it helps memory capacity bound cases (such as large batch) by reducing the working set of DL training.

final_score: rating: 7: Good paper, accept
review_content: Claims:

-Better eviction heuristics provide noticeable improvement over earlier dynamic rematerialization schemes - up to 50-80% memory savings at the cost of up to 75% increase in computation

-Seamless implementation with PyTorch; user does not need to change their code - highly impactful if incorporated into PyTorch

-Upper bound analysis shows only O(N) operations are needed (constant factor more than without rematerialization)

-Several eviction heuristics are studied, including from previous literature. Mathematical formalism unifies these heuristics and clearly describes their relationships

Pros:

-It seems that previous rematerialization schemes provide up to 50% memory savings (i.e. can run 2x larger models), but the proposed work goes up to 80% memory savings (i.e. can run 5x larger models). If this is true, then the ability to run 5x larger models on the same hardware shifts rematerialization from a relatively minor optimization to a game-changing feature

-Running time overhead below 100%, which is an acceptable tradeoff (authors consider runtime overhead >100% to be thrashing, which is sensible)

-More so than typical papers, the design motivations are factually and intuitively explained, and performance claims are not exaggerated but rather put into perspective. The writing strikes me as convincing.

-Comprehensive experiment design with 8 models and 7 eviction heuristics, covering RNNs, CNNs, Residual networks, Transformers, and unconventional models such as Unrolled GAN. This leaves little doubt as to the proposed method's effectiveness

Cons:

-The paper did not study rematerialization in at least the multi-GPU, if not distributed-parallel setting. I think the latter can be excused given the typical scope of an ICLR paper, but I was expecting the former and was surprised to find the experiment design only involved one GPU. At minimum, a multi-GPU experiment is needed to confirm the implementation works with multi-GPU setups, if not distributed-parallel setups.

Questions:

-The paper does not talk about datasets used. This is appropriate if all models in the experiment suite do not have dynamic structure, since dynamic structures could (depending on the implementation) cause the computational graph to differ with different datasets - e.g. I imagine this could happen with NLP and BiLSTM or recursive neural networks. Can the authors clarify if they limited their models to those with static (fixed computational graph) structures? How would the proposed method behave on models with dynamic structures that change with the input data?

Suggestions for improvement:

-The authors call out the relationship between rematerialization and tensor swapping only in the related work. While this is better than not bringing it up at all, it can be argued that since rematerialization and swapping are both dynamic memory management techniques, swapping should at least be mentioned in the introduction. Even if swapping is out of scope of the paper, I think some discussion on how DTR and swapping could be combined would make the paper even stronger.

-Although the paper does explain and provide experiments as to why dynamic rematerialization is superior to static rematerialization, I think the argument could be made much more convincing if the authors showed evidence that static rematerialization does not work on certain dynamic models, limiting its applicability. For example, a statement to the effect of e.g. "Checkmate doesn't work on model X, and here is why..." could be instructive.

review_title: The paper proposes an implementation of tensor rematerialization for PyTorch, in order to reduce GPU memory requirements by up to 50-80% depending on model. The implementation is dynamic (does not require pre-runtime analysis) but requires between 0-75% overhead during runtime. Several heuristics are proposed and analyzed to minimize overheads while maximizing memory saved.

Meta-review: The paper presents an online algorithm for dynamic tensor rematerialization.  The theoretic analysis on the tensor operation and memory budget bound of the proposed method, as well as on the relationship between the proposed method and optimal static analysis method is novel and interesting.  It covers a pretty comprehensive study across theory, simulation and system implementation. In addition, the paper is well written. 
Decision: Accept (Spotlight)