Paper: H-Rockmate: Hierarchical Approach for Efficient Re-materialization of Large Neural Networks (ICLR-2024)

Reviewer comment (Theory):
The partition algorithms (especially the score function/cost model in equation 1) are a bit ad-hoc. While the intuition behind it can be grasped—for example, it tries to identify a subgraph with least IO and penalize the number of nodes in it to minimize the memory required to checkpoint its IO while keeping the scale of each subgraph relatively small—graph partition is a long-studied problem. Usually, such heuristic/greedy-based algorithms don’t scale very well in the sense that they are typically tailored for specific known targets and would fail over time when the target evolves.

Author response:
Thank you for your comments.
* About the partitioning algorithm: please, see our comment on graph partitioning in the General answer. We agree completely that this partition algorithm is ad hoc. We are aware that graph partition is a long-studied problem, but it is not clear what should be the optimization objective for our case. Despite not insisting on it in the paper, the partition algorithm *is* indeed modular in the current implementation and it is just as easy to use another partition algorithm as it is to use another solver. 
  As you mentioned, this paper's contribution is focused on the idea of hierarchical decomposition and the H-ILP formulation. The partitioning algorithm is provided for completeness of the paper, and its purpose is to show that there exists at least one partitioning algorithm that makes this hierarchical idea work. Trying and comparing other partitioning approaches is clearly part of our future work, to see if it is possible to make the idea work even better. And thank you for the suggestion, we will add Densenet to a set of architectures we experiment with.

Reviewer comment (Experiments):
It is suggested to try the partition algorithms on a DenseNet to see if they produce good results.

Author response:
Thank you for your comments.
* About the partitioning algorithm: please, see our comment on graph partitioning in the General answer. We agree completely that this partition algorithm is ad hoc. We are aware that graph partition is a long-studied problem, but it is not clear what should be the optimization objective for our case. Despite not insisting on it in the paper, the partition algorithm *is* indeed modular in the current implementation and it is just as easy to use another partition algorithm as it is to use another solver. 
  As you mentioned, this paper's contribution is focused on the idea of hierarchical decomposition and the H-ILP formulation. The partitioning algorithm is provided for completeness of the paper, and its purpose is to show that there exists at least one partitioning algorithm that makes this hierarchical idea work. Trying and comparing other partitioning approaches is clearly part of our future work, to see if it is possible to make the idea work even better. And thank you for the suggestion, we will add Densenet to a set of architectures we experiment with.

Reviewer comment (Reproducibility):
Alternatively, it is suggested to make this partition algorithm modularized as the solvers. The more valuable/solid aspect in this work is the intuition of hierarchy (divide and conquer) and the H-ILP solver.

Author response:
Thank you for your comments.
* About the partitioning algorithm: please, see our comment on graph partitioning in the General answer. We agree completely that this partition algorithm is ad hoc. We are aware that graph partition is a long-studied problem, but it is not clear what should be the optimization objective for our case. Despite not insisting on it in the paper, the partition algorithm *is* indeed modular in the current implementation and it is just as easy to use another partition algorithm as it is to use another solver. 
  As you mentioned, this paper's contribution is focused on the idea of hierarchical decomposition and the H-ILP formulation. The partitioning algorithm is provided for completeness of the paper, and its purpose is to show that there exists at least one partitioning algorithm that makes this hierarchical idea work. Trying and comparing other partitioning approaches is clearly part of our future work, to see if it is possible to make the idea work even better. And thank you for the suggestion, we will add Densenet to a set of architectures we experiment with.

Reviewer comment (Writing):
The presentation can be improved; for example, the 3.1 H-partition part contains many implementation details without much explanation of where they come from (as mentioned above regarding equation (1)), including why alpha is needed and why it is 0.5, and why 4 candidate groups were chosen in "Formation of candidate groups" rather than another number.

Author response:
Thank you for your comments.
* About the values of the parameters: the value of the parameter alpha of the algorithm was chosen empirically from a small set of experiments, the value $0.5$ seemed to provide the best results overall. About the 4 candidate groups, this is not really a parameter: for a given node $x$, it is not clear whether the "right" partitioning around $x$ is to group it with its predecessor nodes, or with its successor nodes. This is why the algorithm creates one candidate group where $x$ is within its group and one candidate group where $x$ is not within the group. The same goes for $a(x)$, and this yields 4 $(=2\times 2$) possible combinations.

Reviewer comment (Presentation):
The caption of the most important figure (Figure 1) does not have enough details. It is unclear what the time vs. memory plot represents (possibly referring to options), and what “direct solver” means (only base solvers for bottom subgraph and H-ILP hierarchical solver are mentioned).

Author response:
Thank you for your comments.
* About the caption of Figure 1: you are correct that the time vs memory plots represent the set of options. A direct solver is a solver that does not use the hierarchical decomposition but works directly on the whole subgraph itself flattening the hierarchy and all the sub-sub-graphs (TW-Remat is an example of a direct solver). We will make this more clear in the caption of Figure 1 in an updated version.

Reviewer comment (Evaluation):
It is unclear how H-ROCKMATE beats the baselines in Unet/Encoder-Decoders more concretely. An explanation or preferably an illustrative example would help readers understand the quality of it more intuitively.

Author response:
Thank you for your comments.
* About how H-Rockmate beats the baselines: in the example of U-Net, the graph partitioning detects the main structure of U-Net. It is decomposed into a hierarchy where each subgraph corresponds to a layer and contains about 10-15 nodes (see Table 3 in Appendix). This allows H-Rockmate to see the whole structure (which Rockmate can not do) while keeping the size small enough (which Checkmate can not do). TW-Remat is only a heuristic and it is more difficult to guarantee the quality of its schedules.
    In an updated version, we will add (in the Appendix) an illustration of how the algorithm works on the U-Net network, with the hierarchical decomposition and the resulting optimized sequence.

Reviewer comment (Experiments):
Since they claim H-Rockmate works for large neural networks, the sizes of neural networks used in experiments are the same as other works.

Author response:
Thank you for your comments.
* Regarding the size of neural networks: please, see our General answer to all reviewers, where we refer to different definitions of "large" models and specify in which scenarios H-Rockmate is applied.  We have tested much larger neural networks (in terms of number of nodes) than what was previously accessible to Checkmate/Rockmate and obtained performance better than existing approaches. Applying directly H-Rockmate to neural networks whose parameters don't fit one device, like LLaMa, is not feasible because H-Rockmate only optimizes the memory used by activations. However, as we say in the General answer, H-Rockmate can be applied separately to the submodule of the neural network. And as most solutions for training such large models *do* use a form of re-materialization, and our work provides an improved and generic approach for optimizing the re-materialization strategy, we see a task of finding optimal combination of pipeline training and generic re-materialization as future research.

Reviewer comment (Experiments):
If H-Rockmate is applied to a billion-level neural networks like LLaMA, how will the peak memory and iteration time be? I think experiments with larger neural networks than GPT2 is necessary.

Author response:
Thank you for your comments.
* Regarding the size of neural networks: please, see our General answer to all reviewers, where we refer to different definitions of "large" models and specify in which scenarios H-Rockmate is applied.  We have tested much larger neural networks (in terms of number of nodes) than what was previously accessible to Checkmate/Rockmate and obtained performance better than existing approaches. Applying directly H-Rockmate to neural networks whose parameters don't fit one device, like LLaMa, is not feasible because H-Rockmate only optimizes the memory used by activations. However, as we say in the General answer, H-Rockmate can be applied separately to the submodule of the neural network. And as most solutions for training such large models *do* use a form of re-materialization, and our work provides an improved and generic approach for optimizing the re-materialization strategy, we see a task of finding optimal combination of pipeline training and generic re-materialization as future research.

Reviewer comment (Experiments):
What if modeling on-chip global memory to get a better scheduling? Can your method support this?

Author response:
Thank you for your comments.
* Answering Question 3: We are not sure we understand fully what "on-chip global memory" refers to. We assume that it relates to some memory external to the GPU, either the memory attached to the CPU or the memory of one/several other GPUs. This is related to "offloading" problem, where some activations or weights are moved to and from other memory places which incurs a communication delay. Supporting offloading is one of our future work directions, we, particularly,  mention this in the Conclusion and General answer.

Reviewer comment (Evaluation):
I am a little confused by the presented results in Figure 3; I was expecting that when the budget is very low, every algorithm should be able to find the scheduling of re-computing every activation, while when the budget is very high, every algorithm should be able to find the scheduling of no-recomputation, but still there is some difference between each line. This was confusing. Additionally, the important hyper-parameters, such as batch size, are not enumerated in Section 4.

Author response:
* Regarding the results in Figure 3: 
  * About the **low budget** case, your intuition is indeed correct for the case of chain networks, with homogeneous activation sizes. Take for example the case of a chain of length 4. The recompute-all schedule would be F1-F2-F3-F4-B4 - F1-F2-F3-B3 - F1-F2-B2- F1-B1 (where F2 and B2 state for the forward and backward pass through the second layer, correspondingly). If all layers and activations are the same, the peak memory is reached during each B operation. To achieve the minimum budget, no activation can be stored in memory during any B operation: only the recompute-all schedule achieves this minimum budget.
    **However**, when sizes are heterogeneous, it is possible to reach the same memory as the “re-compute all” schedule, but with a lower running time: the peak memory is achieved for the most memory-demanding operation (for example the B3 operation). All other operations can be performed with some activations stored in memory: we can maybe store the result of F1 during the computation of B4 without increasing the peak memory. This allows for a shorter schedule that does not recompute everything. Finding such a shortest schedule still leaves room for optimization. 
	**Furthermore**, when the dependency graph of the neural network is not a chain, other optimization opportunities arise. Take for example a graph where the output of the first layer 1 is used by two paths, 2 -> 3 and 4 -> 5, with a last layer 6 using the result of 3 and 5. The graph looks like
	```    
	       2 -> 3
	      /      \
         1        6
	      \      /
	       4 -> 5
	```
	In a recompute-all schedule, one wants to compute F6B6 with no other activation in memory. Computing F6 requires the result of F3 and F5. While computing F3, one can choose whether to keep the result of F1 in memory or not, and this decision has an impact on the duration of the schedule. For more complex graphs, this decision is not trivial. It was actually shown in [Naumann, 2018] that it is NP-hard to decide whether a given task graph can be computed within a given budget.
  * About the **high budget** case: the best solution is indeed to store everything and to perform no recomputation. However, the values presented in Figure 3 represent actual measurements, as opposed to the estimated duration of the computed schedules. In that case, the ILP-based algorithms do output the "no-recomputation" schedule, but the real-life execution contains overhead from the Rockmate framework.
    Your comment made us reconsider Figure 3b specifically because indeed there should not be such a difference in overhead. We experimented again and obtained a very similar plot, but this time all algorithms obtained the same performance for the high-budget case. The performance is just slightly lower than *Autodiff* because of the overhead of the framework.
    
  * For the **TW-Remat** solution specifically, this one is a heuristic and has no guarantee of finding these good solutions ("recompute-all", "no-recomputation") when they exist.
* Regarding batch size and sequence length: we will mention them in an updated version of the paper. However, the reason why these values are not mentioned is that they only scale the plots (both in terms of memory usage and iteration time), as shown in Figures 6, 7, and 8 of the Appendix. In the experiments of the main paper, the values were selected so that the memory usage is consistent with the total memory available.