H-Rockmate: Hierarchical Approach for Efficient Re-materialization of Large Neural Networks

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.
• 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.
• 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.
• 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.

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.
• Answering Question 2: even without the technical details, the list of H-ILP constraints is significantly long, and highly redundant with the Rockmate paper. Since we wanted to keep the paper easy to understand, we decided to focus on the specific contributions that make H-ILP able to work in a hierarchical manner.
• 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.

Thank you for your comments.

• Regarding the writing of the paper: the first paragraph of Section 3 is our statement of the scheduling problem. It is somewhat short of the necessity of keeping the paper length within the page limit. We will provide more details in an updated version, especially about the input data (execution time and memory overhead of each node in the graph, memory size of each activation). In Section 3.2, as discussed in the main answer, we have indeed opted for an informal description of the specificities of H-ILP compared to the ILP in Rockmate. This was again motivated by the page limit; a more technical presentation would require us to introduce all variables and notations and we could not identify what to remove from the paper in compensation.
• 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.