Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper investigates the advantages of _combining_ the memory optimizations of rematerialization and offloading for training large models on single GPUs.  In doing so it studies when such a combination is useful and unsurprisingly shows that a combination of the two provides better benefits than either optimization used in isolation while incurring a modest 20% runtime overhead.
main_review: I enjoyed reading your paper and your focus on memory saving optimizations for activations to push the boundaries of single GPU training.  I was however quite confused by all the discussion on offloading weights (lines 109-115, references to deepspeed and such).  My suggestion: just explicitly state up-front that your focus is on activations (and you assume weights and optimizer state fit in memory) and just ignore all this distracting discussion on offloading of weights.

I was a little surprised that when it came time to evaluate you picked a batch size of 16.  Could you provide a rationale as to why you pick this batch size for all the DNNs you evaluate with?  Is this something that you can fit for _sequential_ for each of the 8 DNNs you evaluate with? I assume you do this so that you can measure overhead with respect to sequential.  And if so, when you throw your memory saving optimizations at it, are you leaving memory on the GPU unused?  It is ok if you are doing this just for evaluation, but it would be good to call this out.

Is there a reason why your graph x-axis (alpha) ends at 0.8?  Why not show us what happens at 0.9 and 0.95?

One thing that seemed like a missed opportunity in studying this combination of optimizations was reporting what percentage of activations were offloaded versus what percentage were recomputed.  I suspect that because offload is limited by PCIe throughput, offload is used sparingly and spread out by large temporal intervals so that prefetching these tensors can be hidden.  It would be great if you can provide some quantitative data and intuition for this.

As I am writing this review, I realized that I am assuming you are prefetching tensors that you had offloaded (and that you are not doing this fetch in the critical path).  When you offload, are you assuming that you have space buffers available to prefetch data into?  That is, from the total memory available to you are reserving some memory for prefetch buffers?  And is offloading an option only because you are using rematerialization in between - that is you are increasing compute time and hence providing more of an opportunity to hide prefetching from slower CPU memory over PCIe (12 GB/s)?

Overall, while you show us that the combination of the two memory optimizations you consider are effective (compared to individual optimizations), it is unclear _why_ offloading is helping in this combination and how much offloading is being performed.  You have an opportunity to nail this analysis (even if it means suggesting offloading is used sparingly but should be used nonetheless in this combination).

limitations_and_societal_impact: Limitations are mostly covered in my review above, but succinctly: the paper could do a better job explaining how offloading is used in conjunction with rematerialization - that is not just the effect of using them, but what is the sequence of offload/rematerialization operations are the algorithms in the paper suggesting (some sort of temporal view for intuition or percentage of tensors offloaded versus rematerialized for quantitative data).
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: The paper tackles the memory bottlenecks of by activations that limits scaling of model and batch sizes in DL training. The approach is to discover the optimal combination of activation re-materialization and offloading to CPU memory that will minimize the runtime overhead. The optimal solution is found using a proposed dynamic programming algorithm, called `pofo`, which determines an optimal sequence of re-materialization and offload decisions for the network layers, under the constraint of GPU memory and communication bandwidth to CPU. The paper also presents simulation-based results showing that `pofo` more efficiently reduces activation memory pressure compared to different SOTA approaches for different image models and different memory and communication bandwidth constraints.
main_review: This paper tackles a very important problem because GPU memory is perhaps the biggest bottleneck in DL since it grows at significantly slower rates compared to model, data, and compute.  Also, addressing activation memory, which is one of the major sources of memory consumption, is very impactful to the community. A solution that consists of optimally combining SOTA techniques is very appealing and the automated nature could realistically save the massive wastage of human and machine resources on tuning SOTA techniques. 

The paper is very well written and I found it enjoyable to read end-to-end. I felt the authors did such a great job in how they used available space fit enough high-level intuition and low-level details such that I could fully grasp the paper without referring to the appendix.

Concerns:
1. A key weakness is that evaluation does not consider multi-gpu scenarios, which are more common for large-model training than single gpu. This seems like an unnecessary omission that is not adequately explained. Moreover, since your results are simulation-based it does not seem like much extra work is required. It seems that all that is needed is to consider PCI bandwidths per GPU lower than 12GB/sec (e.g., 3GB/sec per GPU for 16 GPUs in a DGX-2 box). This way you won't need to adjust your algorithm to consider inter-gpu communication (which is typically fast within a node anyways). Please clarify if I am missing anything. 
2. I can't figure out if your results apply to language models. The key difference of language models that could affect the effectiveness of your approach is that model and optimizer consume much more memory leaving very little room for activations. On the other hand, and perhaps as a consequence, the batch size is typically much smaller than image models. So perhaps these two factors will cancel out and your results will still hold, but it is hard to tell since you don't have language model results. Another view is that perhaps your work is more desperately needed in the language model space where batch sizes are heavily constrained by memory pressure. Can you share any thoughts on this? Do the claims in the paper need adjusting to reflect the image-centric nature of the results?

Suggestions/questions:
1. The evaluation should include model and activation sizes, as well as the GPU memory size for the presented results.
2. Offloading to NVMe is already possible using ZeRO-Infinity and Nvidia GPUDirect (either from CPU memory or directly from GPU memory). This would be needed when CPU memory is not large enough for the activations. Can your approach can be extended for this?      
limitations_and_societal_impact: Yes.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This paper tackles the problem of combining rematerialization with offloading by using dynamic programming, and demonstrate (in simulations) that their approach can substantially outperform pure offloading as well as pure rematerialization in terms of runtime in different situations.
main_review: Overall, I like the problem that this paper tackles - how do you balance various strategies for making your model bigger/faster? All of these different approaches have different tradeoffs, and it's often unclear how you should use one vs. the other. So, papers that attempt to resolve these questions are quite valuable.

The DP itself is fairly straightforward but reasonable.

My primary qualm with this paper is that the experiments are all performed with a simulator. In my opinion, the DP in this paper is not novel enough to have intrinsic novelty. However, the DP algorithm can still be a valuable contribution if they demonstrate that the DP actually works in practice. Much of the contribution of DP algorithms in situations like this is demonstrating that the assumptions made by the algorithm are realistic enough for the DP to result in a good solution. However, if you're simply running it on a simulator, then I'm not sure this successfully tests whether your algorithm's assumptions are realistic.

Regardless, I think that the simulator results are quite interesting (with good baselines), and demonstrate that offloading may be a technique that people trying to reduce activation memory should consider looking into. 

To be honest, I am quite split on this paper. On the one hand, I think that this paper tackles a valuable problem, demonstrates some novelty in combining rematerialization + offloading, and has intriguing simulator results. On the other hand, without experiments on real hardware or actual software, I am skeptical that these paper's results will have much impact in the community.
limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper presents an algorithm, called pofo, for optimizing the execution time of DNN training on a single GPU, considering a combination of rematerialization and offloading strategies.
main_review: Major Comments
1.	Rematerialization and offloading have been already proposed in the literature. It is not clear which are the challenges of combining both approaches, if any. Therefore, the contributions do not seem much novel.
2.	In practice, can the proposed framework adapt to a limited memory size as a given constraint?
3.	Does the proposed method work only with feed-forward networks, or also with recurrent networks?
4.	In Section 5.2, “In some cases autocapper is marginally better than pofo, as pofo produces solutions under assumptions of Section 3.2, while partial memory releases (assumption (iii)) are difficult to implement in practice.” Why are there different assumptions for the autocapper procedure? For comparison purposes, the assumptions should be the same for every algorithm. Moreover, the reasons why the autocapper outperforms the proposed pofo should be discussed more comprehensively.

Minor Comments
1.	There is no need to report Figure 1 and Table 1 in the appendix, since they are already reported in the main manuscript.

limitations_and_societal_impact: Yes
review_title: null

Meta-review: This paper proposes an optimized algorithm to compute a sequence of forward / backward / offload / prefetch operations on activations that optimizes training throughput of linearized DNNs under memory constraints. The rebuttal solves the reviewers' concerns, and the reviewers unanimously agrees to accept the paper. 
Decision: Accept (Poster)