A fast heuristic to optimize time-space tradeoff for large models
Authors: Akifumi Imanishi
Zijian Xu
Masayuki Takagi
Sixue Wang
Emilio Castillo
Venue: NeurIPS 2023 poster (2023)
Keywords: Recomputation
Gradient checkpointing
Memory reduction
Computational graph optimization
TL;DR: A heuristic that runs in seconds and can reduce memory required for training by 73% with extra 18% computation on average

Abstract: Training large-scale neural networks is heavily constrained by GPU memory. In order to circumvent this limitation, gradient checkpointing, or recomputation is a powerful technique. There is active research in this area with methods such as Checkmake or Moccasin. However, both Checkmate and Moccasin rely on mixed integer linear programming or constraint programming, resulting in limited scalability due to their exponentially large search space.

This paper proposes a novel algorithm for recomputation (FastSA) based on a simulated annealing heuristic that achieves comparable or even better solutions than state-of-the-art alternatives. FastSA can optimize computational graphs with thousands of nodes within 3 to 30 seconds, several orders of magnitude faster than current solutions.

We applied FastSA to PyTorch models and verified its effectiveness through popular large vision and text models, including recent language models with the transformer architecture. The results demonstrate significant memory reductions by 73% with extra 18% computational overheads on average. Our experiments demonstrate the practicality and effectiveness of our recomputation algorithm, further highlighting its potential for wide application in various deep learning domains.

A fast heuristic to optimize time-space tradeoff for large models
Akifumi Imanishi Preferred Networks seanissh@preeeree jp
Zijian Xu
Masayuki Takagi Preferred Networks
Preferred Networks joe@preferred. jp
atakagi@preefrreen
Sixue Wang Preferred Networks cecilvang@preferred.
Emilio Castillo
Preferred Networks
ecastill@preferred
Abstract
Training large-scale neural networks is heavily constrained by GPU memory. In order to circumvent this limitation, gradient checkpointing, or recomputation powerful technique. There is active research in this area with methods such as Checkmake [9] or Moccasin [3]. However, both Checkmate and Moccasin rely mixed integer linear programming constraint programming, resulting in limited scalability due to their exponentially large search space.
This paper proposes a novel algorithm for recomputation (FastSA) based on a simulated annealing heuristic that achieves comparable than state-of-the-arr alternatives. FastSA can optimize computational graphs with thousands of nodes within 3 30 seconds, several orders of magnitude faster than current solutions.
better solutions
We applied FastSA to PyTorch models and verified its effectiveness through pop- ular large vision and text models, including recent language models with the transformer architecture. The results demonstrate significant memory reductions by 73% with extra 18% computational overheads on average. Our experiments demonstrate the practicality and effectiveness of our recomputation algorithm, further highlighting potential for wide application in various deep learning mains.
1 Introduction
The memory requirements for deep neural networks continue grow together with model complex- ity. For example, training state-of-the-an Large Language Model (LLM) such as LLaMA [4) with 65 billion parameters needs 1024 NVIDIA A100 GPU devices for 21 days and different tech- niques to overcome the immense memory consumption of the process. While the model parameters may fit in the device memory, the results of several operations of the forward pass are saved and remain in the device memory to compute the gradients during the backpropagation step. This limits size of trainable models. One of the most widely used solutions is to split the model parameters across multiple devices into vertical or horizontal ways 34 26 in what is called Pipeline or Model parallelism. However, these approaches may have the following problems. (I) may be difficult to split the model equally across the workers when the model architecture is complex. (2) We need to modify the model for the parallelism. (3) Hiding communication behind computation needs fur- ther tuning. Heavy inter-device communication may result in a GPU utilization of only 5% of the hardware peak[3] 30
"Equal contribution
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Optimization time [5]
10
TestSA (curs)
Checkmate LP 25% budget
107-
50% budget
10
500
1000
2000
4000
0200
Number of nodes
Figure 1: Optimization time for all 23 models used in our experiments with different budgets. Checkmate LP could not find feasible solution within 6 hours for models with more than 2300 nodes. For most instances, FastSA A could found recomputation plan around 3 seconds.
Numerous techniques exist lessen memory requirements while training deep neural networks. Gradient checkpointing, or recomputation, is widely-knows approach that recomputes some of the activations during the backward pass, rather than saving them memory, consequently reduc- ing memory requirements. Recomputation plans can be identified by the user or done heuristically. Originally, division-based methods were used for automatic recomputation, where the forward com- putation graph was split into stages, and only tensors active across stages were stored 7 21 20) However, these methods are less effective for complex large networks that are far from sequential networks. Checkmate is of the state-of-the-art methods in automatic recomputation. It can be applied to any computational graph, considering the costs of operators and the sizes of values. However, Checkmate requires to solve large mixed integer linear programming problem, where the search space is exponential to the square of the size of the computational graph, thus requires immense computational time and large RAM resource. Moreover, Checkmate cannot efficiently reduce the memory usage if the initial computation order is not memory efficient enough (discussed in Appendix A. A recent work, Moccasin [3), uses constraint programming to build set of re- strictions and introduces hyperparameter that limits the amount of times a variable can be recom- puted. Although the search space remains exponential like Checkmate, this hyperparameter lessens the number of integer variables from quadratic to linear with respect to the graph size, resulting in a faster execution.
This paper proposes a novel technique to recompute the activations that need to be saved for the backpropagation pass by using fast heuristic algorithm that operates the joint computational graph of the forward and backward passes and determines which operations to be recomputed to reduce the amount of used temporal memory. We formalize the problem of recomputation as find- ing a candidate sequence of operators that minimizes an objective function. We apply a simulated annealing algorithm where each step be operator addition deletion from the computational sequence, or altering the order in which operators are sequentially applied, while maintaining the dependencies between them. The performance of the simulated annealing significantly improved by using a segment tree data structure [9] which allows to lazily evaluate the memory usage of each candidate sequence, enabling nearly million mutations the candidate sequence per second.
The main contributions of our paper are:
We present a novel approach to recompute intermediate results in computational graphs using heuristic algorithm based on simulated annealing that works for any computational graph.
We further improve the algorithm's performance using a data structure that allows efficient computation of peak memory usage when the sequence of operators is changed. This can optimize the recomputation within 3 to 30 seconds even for large computational graphs.
We evaluate our proposal using a representative set of models obtained from Hugging Face, including vision models, text models, and recent language models such as LLaMa [40] We show reductions in memory consumption by 73% with an overhead of extra 18% compu- tational cost on average.

---

time
time
Graph
Graph
value recomputed time
Lifetime
Lifetime
Figure 2: Lifetime of recomputed values. The left side shows the full lifetime of a value without recomputation, while the right side demonstrates that recomputation results in the value having two distinct lifetimes.
Fast Simulated Annealing for Recomputation
2.1 Prohlem Description
To apply the Simulated Annealing optimization strategy the recomputation problem, we present a series of formal concepts that will be used throughout the algorithm description.
Computational Graph and Sequence We consider computational graphs in which each node represents an operator. A computational graph consists of finite values V. Each node takes some values as inputs and produces some values outputs; therefore, we introduce notations input(n) S V and output(n) In most cases, the numbers of (n) and output(n) are small. The computational graph is often visualized as a directed acyclic graph (DAG) with directed edges from nı, E N to E output of n, is used as input of n.j i.e., output(n:: input(n;) & 0. A value that is not the model input, which often represents the arguments passed to the model and the model parameters. We also have model outputs, the target of the computation. Both nodes and values are weighted: the computational cost of node n is represented by c(n) 2 0, and the size of value e is represented by s(u) 2 0.
nodes N and finite set of
V for the values associated to node n W.
an
output of any node
Rather than performing the optimization over the DAG, the proposed method is easier to understand if the computation is treated as a sequence of nodes (M,,..., Rr), which means that is operated integer time We say that sequence of nodes has valid dependency if we can execute the operators in the order of the sequence so that all model outputs are correctly computed. Formally, Vi.Vu E input(n:).33 < such that E output(n;) or v is a model input. It is worth noting that node duplication is allowed in the above node sequence (1.,,.,,T))) and the recomputation of value can be represented by multiple appearances the same node.
Memory usage In order to optimize memory usage, we must keep track of when values are com- puted and used inputs. Specifically, as an output until the last time is used as an input. However, when v is recomputed at time f, we temporarily free from memory if it is not used until t. When recomputation is considered, the lifetime of v can be represented as set of independent time intervals, as shown in Figure 2 To determine whether v is kept in memory given time I, we function Ls(v, {0, 1). This function outputs if v is kept in memory at time I, and 0 otherwise.
must keep value in memory from the time it is computed
{
1 u is model input u a model output A VE' > u output(ne))
Ls(u,t)
3t" > t input(ne:) A Vt' [t,tt), & output(ne) otherwise
Let MS) be the minimum memory size required to compute a valid sequence S. Then, we have M(S) maxigIsT Evev Ls(v,t) s(v). Also, we define the cost of the sequence, C(S), by C(S) = Eisisr c(n)).

---

Choose new state Add node
Mernory Simulation Data Structure
Update State
Memory
Gel Peak
Optimized
Enough?
Yes
3.
Swap nade
of New State
No
End
Initill
Set New
Revert State
No
Accept
New
Temoerature
State?
Yes
Figure Overview of the simulated annealing.
Recomputation Let f be an objective function which maps sequences to a real number. The re- computation problem is defined as the construction of valid sequence (M1,...,,yy) nr) that minimizes f(S).
When minimizing memory usage, set S(S) MS). The objective function can depend not only on (S) but also on computational costs. our experiments, we utilized Equation
f(8) max(budget, M(S) C(S)
(1)
This function minimizes the total computational cost when the memory budget is met. Otherwise, tries to minimize both the peak memory and the cost.
2.1.1 Add-max segment tree
The uugg-adddrange-man max segment tree [ is a data structure that can evaluate add and sum queries for intervals lazily. The add-max segment tree, utilized in this paper, holds array A[T] of length T, and can carry out the following operations in O(log T) time:
Add A[i] for i (,,r)
Obtain the maximum value of A[i] for i E [,,r)
2.2 Fast Simulated Annealing Algorithm
The proposed algorithm is based on simulated annealing (SA) with straightforward outline. We make small modifications to the sequence and accept them if they decrease the objective function f(S). Please see Figure 3 for overview of the algorithm. Two critical mutations must be per- formed in sublinear time for the algorithm to be efficient: (1) Slightly modifying S to create S" (2) Computing fS").
To generate from S, we either add or remove a node. We represent the sequence as a fixed- sized vector of nodes, where each position in the sequence corresponds specific time step. This representation makes it easy to manage the lifetime of each value. The vector contains a special node nop, which has zero computational cost and inputs outputs associated with it. When adding or removing node, its value is exchanged with that of nop.
We start with a valid sequence S = (nI,..., IT) and for all 1 SiST. E (N U ). To enable insertion of a long computation between nodes, most elements of S should be initialized with nop. In practice, the initial S is constructed from the sequence representing the initial execution order by inserting nop nodes between each element.
In each iteration, we perform one of the following three mutations:
1. (Add computation) Select random node n A and random time t where n, nop. We attempt to update i.e., f-
2. (Remove computation) Select a random time / such that n4 * nop and attempt to update M4 nop.
3. (Rotating sequence) Select two random times ( and l2 such that n4, * nop and nı, nop. Try swapping n,, and nty- This modification can implemented with the combination of the first two modifications.

---

To ensure that the mutated sequence S' has a valid dependency, we validate the following conditions and update the sequence only they are satisfied. Note that we can efficiently check these conditions by maintaining set of produced and used times for each value v.
1. (Add computation) For each imput of ne, must be produced before time t. 2. (Remove computation) For each output of ny, either must already be produced before time i, or there must not be any user of e before the next production of nı:-
An important point note is that the mutations [ are inverses of each other, meaning that to undo change, one simply needs to reapply the inverse mutation.
2.2.1 Updating peak memory and objective
In the following discussion, consider the update of f(S) by mutation I and 2 Since the objective function f(S) often depends on both memory usage (S) and the overall cost C(S) = E, c(n,), assume this is still the case here. We can easily update C(S) by adding corresponding cost of a newly added or removed node from S. Updating M(S) is more challenging; to calculate the maximum peak memory consumption when node is added or removed from S, have to calculate the aggregated memory consumption by the values with overlapping lifetimes.
subtracting the
The lifetime of inputs and outputs change when adding or removing node n. (see Figure S in Appendix for illustration). However, the insertion and deletion of node only slightly modify the life intervals for each node's input and output values. Consider value v with life intervals L(v) := tt) determined by (1, t). Suppose node n is inserted at time t. and v is an input of n. We update the life intervals of as follows: If there is (t,, ty) L(v) such that S t ty, do nothing. Otherwise, take a (z) E I(u) such that 13 S 8 is maximum (such 13 exists the sequence after insertion is valid). Update L(r)
((;) (z) (t,, Similar rules apply for updating the life intervals of the outputs of the inserted node and the inputs and outputs of the node remove. Specifically, on inserting or deleting node, the update of life intervals for an input or output value applies to one of the following four cases:
1. No update.
2. Extend or shrink a range (tz, (R) to (Lz, ( or (tv,, (R).
3. Split range (4z, (R) into (t) and (Iw,,n) (n)-
4. Merge two ranges (() and ((1,,R) into (lz, (R).
Because the update of live intervals involves adding or subtracting the value size s(v) from certain range for each input or output value (as depicted in Figure 5 in Appendix), we can efficiently maintain memory usage and calculate the peak memory using segment tree. This data structure allows for range-max and range-sum queries to be performed in O(log I') time, as introduced in Section 2.1.1 The segment tree maintains memory usage for time t in internal array A[t]. and peak memory can be determined efficiently by taking the maximum value in the range [0, T") in Ollo T') time. To be specific, we add s(v) for the extended range and subtract s(v) for the removed range when updating the live intervals. The segment tree significantly contributes to the simulated annealing
reduce the time for the differential update of a lifetime interval from the naive O(T) time to O Olog 7') time.
2.2.2 Improving memory reduction by grouping
To determine optimal recomputation sequence, it may be necessary replicate specific node patterns to rematerialize a larger value. However, achieving this using random insertion and deletion of nodes is challenging. To address this issue, we introduce the concept of grouped nodes, which concatenates multiple nodes in a series to enable the recomputation of series of nodes. A grouping node represented by 9, is formed by concatenating nodes, n, and n2, and has the following properties:
<(g) = c(n)) + c(n2)
ig)ut(g) input(n1) U (input(ng)) (aaput(na)oouppu(mm]]] (utput(n11)
output(g) = (output(nı) (input)(())))))))) U Joutput(nz)
5

---

Grouped nodes can be especially useful for sequenced patterns where n1 must be computed imme- diately before n2. as grouping these nodes together can make the sequence more likely converge during simulated annealing. Additionally, we observed that after conducting annealing with grouped nodes, further improvement can be achieved by decomposing the grouped node and performing an- other round of annealing with lower temperatures. In Appendix C.I on the benefits of using node grouping for optimization.
we provide further discussion
2.3 Other Considerations
Conceptually, recomputation enables to optimize the time-space tradeoff without changing the graph semantics, i.c., the output values of the neural network. However, in real models, certain operators with internal states or side effects may produce different output values if recomputed or if their computation order is changed. Our algorithm can handle these cases by simply prohibiting the node addition or removal for these nodes.
An additional way reduce memory usage that can be easily incorporated into recomputation algorithm is offloading. which involves moving tensors from GPU memory to host memory when they are not immediately needed, and then moving them back before performing the dependent oper- ation. Our algorithm can be extended to support offloading, as outlined in Appendix (B.3] Supporting offloading together with recomputation offers the potential for even greater reductions in memory usage, particularly for larger computational graphs
3 Experiments
In this section, we present the results of our experiments with the recomputation algorithm. Our algorithm was integrated into the PyTorch framework, and we used it to optimize the internal com- putational graph of various popular models, including vision and text architectures.
3.1 Configuration
Model and input data Our experiments involved the latest vision models and vision transformers obtained from timm (PyTorch Image Models), as well as text models (including language models) from Hugging Face transformers. To obtain the full computation graph with backpropagation, the up for image classification, and model variants with sequence classification
vision models heads for the text models. The computational graphs were obtained by PyTorch's symbolic tracing. The value sizes are estimated from shapes and data types. Since node cost estimations were not available in symbolic tracing, all nodes were estimated to have unit costs. For memory budgets, we used the 50% and 25% values of the simulated initial peak memory. For additional details on the environment, PyTorch integration, hyperparameters of the models (e.g., batch sizes, sequence length in transformers), and the effect of using simulated costs instead of actual ones, please refer to Appendix D
Objective function and hyperparamters of SA As outlined in Section 2 we utilized objective x C(S) for our SA algorithm unless mentioned otherwise? This objective function seeks to minimize C(S) once the memory budget is met, and until that point, it minimizes both memory and cost. We included cost in the memory minimization objectives address the risk of large numbers of nodes, which can result in slower convergence. This objective function is continuous and independent of the cost or size units, and it works effec- tively for a broad range of practical models.
function of the form f(S) = max(budget, (S)) M
To ensure efficient convergence, we utilized the first SA on grouped nodes for at most 20 million iterations, until the memory budget was met, whichever occurred first. The second SA for fixed million iterations for the purpose of cost reduction as detailed in section 2.2.2
Checkmate To the best of our knowledge, the most powerful recomputation planner currently available is Checkmate [19] We re-implemented Checkmate for PyTorch. However, Checkmate
LIn certain scenarios, this objective function may exceed the budget constraint. For experiments of Table and 2 where memory budgets must be satistied, we employed f(S) = ((Elu(L"((MS)) x CCS) where L is linear function with L(0) budget and L 100) 0, determining the sharpness of Elu.

---

FasiSA Checkmate LP
FastSA Checkmate LP
0.6
17-
memory
Relative cost
1.5
0.4-
1.4
0.3
33-
0.1
1.0
ConvNett
VIT (smalf)
BERt
GPT2
FastSA AVE
ConyNeXt VIT (small)
Brtclenthieggat
GPT2
FasISA AVE
Figure 4: Comparison of simulated memory decrease and cost increase by recomputation. For each model, the memory budget is bars represent the 25% budget and light bars represent 50% budget). FastSA AVE represents the geometric mean for all 23 models by FastSA. Checkmate LP's average bar is not shown since failed to solve 11 out of 23 model instances.
50% and 25% of the simulated initial peak memory (darker
could not find any feasible solution within 6 hours for all models with more than 500 nodes due the size of the integer linear programming (MILP) to be solved. Therefore, for the comparison with large models, we resort to Checkmate LP, where the MILP was relaxed to linear programming (L.P). For clarification, the original version is referred as Checkmate MII.P. Gurobi [13 is used as
the internal solver for Checkmate MILP together with PDLP solver, provided by OR-Tools (2 for Checkmate LF
Moccassin Another alternative perform recomputation is Moccasin Similarly Check- mate, Moccasin uses constant programming (CP) to solve the recomputation problem. It introduces a new hyperparameter that acts as an upper-bound on how many times value can be recomputed. Thanks to this limitation, the number of integer variables involved in the CP set is linear to the num- ber of nodes in the graph, allowing faster execution time than Checkmate. However, the possible solutions that Moccasin can converge to are heavily limited by this hyperparameter value and the achievable memory reductions can be sub-optimal.
3.2 Memory reduction on large models
Figure 4 provides comparison between FastSA and Checkmate LP in terms of simulated memory reduction and cost increase. We selected the most representative models for this figure, in addition to the geometric mean of all 23 models tested in Appendix D.7] Checkmate LP failed to find solutions for 11 models due to time limits or out-of-memory errors of 100 GiB, so the geometric mean is not reported. Moreover, memory usage was not fully reduced to the budget due to LP relaxation and randomized rounding, which resulted in larger overhead for some models, such as ViT (small); we discuss this behavior in Appendix D.6
For the 50% memory budgets, FastSA algorithm was successful in meeting the target budget for all models, with average increase in model execution time of 7%. However, for the 25% budges, FastSA could not reduce memory to the budget for some models due to either limitations on the model itself or convergence the algorithm reduced memory usage by an average of 73% with an average increase in overhead of 18%. Overall, these results demonstrate the superior performance of FastSA compared to Check- mate LP in reducing memory usage, particularly for larger neural networks.
suboptimal solution due SA heuristic limitations. Despite this,
A comparison with Moccasin was also conducted by applying FastSA the publicly available data in 3]. Tabld presents the results from the original Moccasin study, extended to include outcomes for FastSA. The experiments pertain to graph variants, namely random-layers (RL) and graphs mentioned in 19 (CM). For the RL graphs, FastSA achieved the lowest recomputation overhead
3 'The PDLP solver is used Checkmate LP LP experiments were executed in parallel using a cluster environment. We verified that for the experiments in F the PDLP solver could solve all the instances solved by Gurobi, making it reasonable open source alternative.
licensing costs associated to Gurobi, Checkmate

---

CHECKMATE MILP
MOCCASIN
FASTSA
Graph
(n, m)
B
|
CI
Mem
Time
CI
Mem
Time
CI
Mem
Time
RL 1
100
90
0.8
89.2
18.5
0.8
88.1
9.3
0.0
79.4
11.7
236
80
2.3
79.5
22.7
2.3
79.5
9.5
0.3
79.4
11.5
70
no experiment
no experiment
2.2
77.3
12.0
RL 2
250
90
0.9
90.0
685.1
0.9
89.8
55.0
0.0
77.3
15.1
944
80
time limit exceeded
4.9
80.0
639.5
0.0
72.2
15.0
70
experiment
experiment
2.6
68.9
14.5
RL 3
500
90
time limit exceeded
0.7
90.0
1803.3
0.03
87.4
21.0
2461
80
time limit exceeded
3.4
80.0 1804.8
2.3
78.6
20.9
70
no experiment
experiment
4.8
(73.5)
21.2
RL 4
1000
90
time limit exceeded
0.7
90.0
3612.9
0.4
87.5
36.0
5857
80
time limit exceeded no experiment
3.4
80.0
3611.8
2.5
78.4
36.3
no experiment
7.4
70.0
36.5
CM 1
73
90
0.0
88.4
6.3
0.0
88.4
3.1
0.1
75.3
9.4
149
80
0.1
76.9
5.6
0.1
78.9
3.1
0.1
75.1
9.7
70
no experiment
no experiment
3.0
62.3
9.7
CM 2
353
90
0.1
89.0
434.1
0.2
89.9
65.2
0.2
86.2
10.9
751
80
0.3
79.7
485.3
0.3
80.0
69.3
0.4
76.6
10.9
70
experiment
experiment
0.8
66.8
11.0
The table includes results of Checkmate MILP and Moccasin from 3] extended with FastSA data. The columns В and CI represent memory budget and cost increase percentage, respectively. Alongside original results for 90% and 80% budgets, 70% budget only demonstrating FastSA results is also inserted. For all random-layered (RL) cases, FastSA exhibited the smallest CI, as optimizes topological ordering, reducing memory without materialization.
Table Comparison with Moccasin.
in all instances. Remarkably, for RLI and RL2 cases, FastSA managed decrease memory usage without adding new recomputation nodes, through the optimization of the execution's 's topological ordering.
3.3 Solution optimality
CHECKMATE MILP
FASTSA
Model
(n,, m)
B
CI
Mem
Time
CI
Mem
Time
VGGI1
(69, 119)
90
1.4
87.6
5.1
2.9
87.5
2.4
80
2.9
79.8
1.5
2.9
79.8
2.4
70
infeasible
2.9
(79.8)
2.4
ResNet18
(171, 437)
90
0.6
85.7
47.6
0.6
85.7
2.1
80
1.8
78.6
45.3
2.9
75.0
2.1
70
2.3
67.9
367.3
4.7
66.1
2.1
60
4.1
57.1
1846.0
5.3
55.4
2.2
50
(4.7)
(50.0)
3600
6.4
(51.8)
2.3
Table 2: Comparison between FastSA and Checkmate MILP small models. For the resnet 8 case with 50% budget, the best feasible solution was used because the MILP solver could find the optimal solution within the time limit.
We conducted study on how closely FastSA solutions match with the optimal solution by com- paring its results with those of Checkmate MILP on small models. The results, compiled in Table 2, show that for the model VGG11 under the strictest budget constraint (80%), FastSA located the same plan as Checkmate MILP for recomputation. For ResNet18, except for 50% budget, FastS A managed to cut down memory use under the budget in all cases, but with up to 2.4% more over- head from recomputation compared to Checkmate MIL.P Even though FastSA fell short of finding the optimal solutions for these cases, it up 1000x faster than Checkmate MILP, particularly

---

under strict memory budget scenarios. In certain situations where the computational graph has high topological freedom, Checkmate might present suboptimal recomputation plans (discussed in Ap- pendix A.1 This can be found in the results for RI. graphs in Table | where FastSA could find better solutions than Checkmate MILP or Moccasin.
4
Related Work
4.1 Model Parallelism
One common approach to scaling single neural network that is limited by memory is to partition it into multiple devices using either Model Parallelism or Pipeline Parallelism. Model Parallelism [34] involves horizontally splitting the neural network by distributing the parameters of each layer across multiple GPU devices, while Pipeline Parallelism [26) proposes to vertically split the network by assigning each device several contiguous layers of the model. While these approaches enable deep neural networks to be trained at scale, their performance can be limited by the communication over- head required for each iteration of the model. Rajbhandari et al. [30] found that these approaches can achieve only 5% of the performance peak of V 100 device, highlighting the limitations of these techniques. These limitations have spurred further research into improving neural network scaling, which is discussed in more detail below.
4.2 Recomputation
Recomputation is technique that was first introduced in classical compiler research to minimize the number of required registers, and later on was adapted for use in deep neural networks (DNNs) by Chen et al. 77 as means of reducing memory consumption during training of sequential models. However, this method is limited to sequential graphs and disregards node costs. Kusumoto et al. (21) proposed dynamic programming algorithms for more general computational graphs. Also, Kumar et al. (20) leverages tree decomposition to handle more general graph structures. However, these methods still requires large computational overhead, making them impractical for larger networks.
Jain et al. [19] formulated recomputation as mixed integer linear programming (MILP) problem and proposed Checkmate, solver to find an optimal recomputation plan. Checkmate tries to min- imize the total computational cost under memory budget and dependency constraints. Although Checkmate has shown to significantly outperform existing methods in terms of solution quality, it requires substantial computational resources to solve the MILP. The number of decision variables in the MILP scales quadratically to the graph size and time scales exponentially to it. To address the limitation of Checkmate, Bartan et al. [3] proposed Moccasin, which formulates recomputation using constraint programming (CP). The number of integer variables is reduced from quadratic to linear in Moccasin, by setting hyperparamter for the maximum times a value can be recomputed, thus expediting execution. However, this formulation also narrows the search space and may impact overall quality of the recomputation plans when compared to Checkmate or A. Also, in a par- allel development, Rockmate 44 was proposed for models with repeated layers. It decomposes the problem of recomputation into intra-layer and inter-layer recomputation. It applies Checkmate to a single layer, which forms smaller graph than the entire model, and then finds the recomputation plan across the layers by Rotor [177 a dynamic programming based recomputation algorithm that works for sequential models.
In contrast, our proposed method converges to solutions that are comparable or even better than Checkmate and Moccasin with single CPU core, taking less than four seconds on average. This significant reduction in computational time is achieved by leveraging efficient heuristics and opti- mization techniques. Our results demonstrate that approach has great potential in the context of real-world applications, especially for cases where large computational resources are not available.
4.3
Other techniques for memory reduction
There are several other techniques that have been proposed to reduce device memory consumption in addition to recomputation, including offloading, which involves transferring some of the model parameters to a system's CPU or an external memory device when they are not immediately re- quired. Beaumont et al. discuss the combination of offloading and recomputation as means of further reducing memory utilization. Our proposed method can easily be extended to support

---

offloading, as detailed in appendix B33 Another approach that leverages offloading is the ZeRO technique proposed by Rajbhandari al. 30]. which partitions the model parameters, optimizer states, and activations among several devices to increase parallelism and reduce the memory used by each device. This approach enables the training of exceptionally large models with hundreds of billions of parameters, making it powerful tool for advanced natural language processing and computer vision applications. Other techniques focus on reducing the size of the parameters and intermediate results, such as quantization [23], which reduces the floating point precision of com- putations and weight storage up to 2-bits extreme cases. Sparsification [25] exploits the sparsity patterns that arise during computation to reduce the total needed memory while keeping the same precision. There are also more exotic approaches, such as the reversible residual network [14), which is memory-efficien architecture that can perform backward computation without saving the activations. However, its applicability is limited to residual networks only. is worth noting that these techniques are largely orthogonal to proposed method and can be combined to further improve memory savings in neural network optimization.
Conclusion
In this paper, we present a novel method for recomputation that offers several key advantages over existing approaches. Our method is applicable general graphs and support any objective function that depends on peak memory usage and the total cost of computation. Moreover, it can find near-optimal recomputation plans within remarkably short computational time of 3 to 30 seconds, even for large computational graphs.
Another major advantage of our method is its efficiency, as it is single-threaded and uses memory resources highly efficient manner. This makes ideal for integration into neural network com- pilers, where it can further streamline the optimization process. Additionally, our algorithm is highly flexible and can handle a wide range problem settings, including recomputation with offloading and consideration of iede-intermmdiite memory usage.
We explored the tradeoff between computation time and memory usage and evaluated the effec- tiveness of our algorithm in terms of reducing peak memory usage. Our experiments demonstrate that our approach can achieve significant reductions in peak memory usage, with reduced impact on computation time, for a wide range of neural network architectures. Overall, our experiments validate the effectiveness and practicality of our recomputation algorithm, highlighting its potential for widespread application in many areas of deep learning research and development.
Limitations In conclusion, while our algorithm offers powerful solution for complex neural network optimization problems. it has certain limitations that must be considered. As demonstrated in Figure 4 our approach was able reduce memory usage more significantly than Checkmate LP in most cases, but the result was not ideal for some models due to suboptimal node grouping.
When planning recomputation in distributed training, the computational graph may need to have extra requirements such as (1) the optimized graph must be able to split equally for each worker in Pipeline Parallel or (2) the communication operators must be overlapped with arithmetic computa- tions. Currently, it is difficult to handle these constraints with the standard FastSA
Additionally, the peak memory estimation may not be always precise due to memory aliasing, caused by operators that change the memory views or do inplace updates. This is one of the reason why there is gap between the simulated and actual GPU memory usage as shown in Appendix D.8
Overall, our proposed approach offers fast and efficient method for neural network optimization, providing significant improvements over existing techniques. However, future research may be required to support wider variety of computation.
References
[1] Alaaeldin Ali, Hugo Touvron, Mathilde Caron, Piotr Bojanowski, Matthijs Douze, Armand Joulin, Ivan Laptev, Natalia Neverova, Gabriel Synnaeve, Jakob Verbeek, et al. Xcit: Cross- covariance image transformers. Advances in neural information processing systems, 34: 20014-20027, 2021.
10

---

[2] Hangbo Bao. Li Dong, Songhao Piao, and Furu Wei. Beit: Bert pre-training of image trans- formers. arXiv preprint arXiv:2106.0825,, 2021.
[3] Burak Bartan, Haoming Li, Harris Teague, Christopher Lott, and Bistra Dilkina. Moc- casin: Efficient tensor rematerialization for neural networks. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, edi- tors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 1826-1837. PMLR, 23-29 Jul 2023. URL
https::/prooceedingg .mr..............m............ mlr wwwpp/ (tress//202/bbrraa2233 html
[4] Olivier Beaumont, Lionel Eyraud-Dubon and Alena Shilova. Efficient combination of materialization and offloading for training dnns. Advances in Neural Information Processing Systems, 34:23844-2355 2021.
[5] JL Bentley. Algorithms for klee's rectangle problems. dept. of computer science, 1977.
[6] Sid Black, Leo Gao, Phil Wang, Connor Leahy, and Stella Biderman. GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh- Tensorflow, March 2021. URL https: If you use this software, please cite it using these
doi wtt.oo//g.//////ee...ee metadata.
5297715
[7] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint 1eXiv:1604006177 2016.
[8] Kevin Clark, Minh-Thang Luong, Quoc V Le, and Christopher Manning. Elec- tra: Pre-training text encoders as discriminators rather than generators. arXiv preprint arXiv:2003. 10555, 2020.
[9] Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein. Introduction to algorithms. MIT press, 2022.
[10] Zihang Dai, Hanxiao Liu, Quoc V Le, and Mingxing Tan. Coatnet: Marrying convolution and attention for all data sizes. Advances in Neural Information Processing Systems, 34:3965- 3977, 2021.
[11] Jacob Devlin, Ming-We: Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
[12] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
[13] Yuxin Fang, Quan Sun, Xinggang Wang, Tiejun Huang, Xinlong Wang, and Yue Cao. Eva-02: visual representation for neon genesis. arXiv preprint arXiv:2303. 11331, 2023.
[14] Aidan Gomez, Mengye Ren, Raquel Urtasun, and Roger Grosse. The reversible resid- ual network: Backpropagation without storing activations. Advances in neural information processing systems, 30, 2017.
[15] Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2023. URL httph:tt/p:/ett //www gurobi com
[16] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for im- age recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[17] Julien Herrmann, Olivier Beaumont, Lionel Eyraud-Dubois, Julien Hermann, Alexis Joly, and Alena Shilova. Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory, arXiv preprint
2019.
11

---

[18] Andrew Howard, Mark Sandler, Grace Chu, Liang-Chich Chich Chen, Bo Chen, Mingxing Tan, Wei- jun Wang, Yukun Zhu, Ruoming Pang, Vijay Vasudevan, et al. Searching for mobilenetv3. In Proceedings of the IEFE/CVF international conference on computer vision, pages 1314-1324, 2019.
[19] Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and Ion Stoica. Checkmate: Breaking the memory wall with optimal tensor rematerialization. Proceedings of Machine Learning and Systems, 2:497-511, 2020.
[20] Ravi Kumar, Manish Purohit, Zoya Svitkina, Erik Vee, and Joshua Wang. Efficient remateri- alization for deep networks. Advances in Neural Information Processing Systems. 32, 2019.
[21] Mitsuru Kusumoto, Takuya Inoue, Gentaro Watanabe, Takuya Akiba, and Masanori Koyama. A graph theoretic framework of recomputation algorithms for memory-efficien backpropaga- tion. Advances in Neural Information Processing Systems, 32, 2019.
[22] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Sorieut. Albert: lite bert for self-supervised learning of language representations. arXiv preprint arXiv:1909. 1942, 2019.
[23] Hao Li, Soham De, Zheng Xu, Christoph Studer, Hanan Samet, and Tom Gold- stein. Training quantized nets: deeper understanding. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https::/pproceeddingg http:/sc.ww.spsppppppppere.rraee cwtpaper_tiiiiiieccee wtt.:/eleeepppe///////////////// pttpp/aapr/211//1le .pf
[24] Zhuang Liu, Hanzi Mao, Chao-Yual Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. A convnet for the 2020s. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11976-11988 1986, 2022.
[25] Dmitry Molchanov, Arsenii Ashukha, and Dmitry Vetrov. Variational dropout sparsifies deep neural networks. Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML' 17, page 2498-2507 JMLR.org, 2017.
[26] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, Gregory R. Ganger, Phillip B. Gibbons, and Matei Zaharia. Pipedream: Generalized pipeline parallelism for dnn training. In Proceedings of the 27th ACM Symposium Operating Systems Principles, SOSP '19 page 1-15, New York, NY, USA, 2019. Association for Computing Machinery. ISBN 9781450368735.
doi: 10.1145/3441301.335944 URL
https://doi. org/10. 10. 1145/3341130 3359646
[27] Wolfgang J Paul and Robert Endre Tarjan. Time-space trade-offs in pebble game. Acta
Informatica, 10(2):111-115, 1978.
[28] Laurent Perron and Vincent Furnon. Or-tools. URL UtL.https::/deevllpprs..go gowgse...ooo.o.oooooooo com/ optimization/
[29] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever, Lan- guage models are unsupervised multitask learners. 2019.
[30] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory op- timizations toward training trillion parameter models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC '20. IEEE Press, 2020. ISBN 978 9781728199986.
[31] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108, 2019.
[32] Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 776b-parameter open-access multilingual language model. arXiv preprint arXiv:2211.051 2022.
12

---

[33] Ravi Sethi. Complete register allocation problems. In Proceedings of the fifth annual ACM symposium on Theory of computing, pages 182-195, 1973.
[34] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penpor Koanan- takool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young. et al. Mesh- tensorflow: Deep learning for supercomputers. In Advances in Neural Information Processing Systems, pages 10435-10444, 2018.
[35] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv:1409.1555 2014.
[36] Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How train your vit? data, augmentation, and regularization in vision trans- formers. arXiv preprint 10ti:2200.102700, 2021.
[37] Quan Sun, Yuxin Fang, Ledell Wu, Xinlong Wang, and Yue Cao. Eva-clip: Improved training techniques for clip scale. arXiv preprint arXiv:2303.1 15389, 2023.
[38] Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In International conference machine learning, pages 6105-6114 14. PMLR, 2019.
[39] Hugo Touvron, Matthie Cord, and Hervé Jégou. Deit iii: Revenge of the vit. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Part XXIV, pages 516-533. Springer, 2022.
[40] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 2023.
[41] Ross Wightman, Hugo Touvron, and Hervé Jégou. Resnet strikes back: An improved training procedure in timm. arXiv preprint 1rXiv:2110.0000
2021.
[42] Sanghyun Woo, Shoubhik Debnath, Ronghang Hu, Xinlei Chen, Zhuang Liu, In So Kweon, and Saining Xie. Convnext v2: Co-designing and scaling convnets with masked autoencoders. arXiv preprint MxXiv:2301.00800 2023.
[43] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained trans- former language models, arXiv preprint arXiv:2205.01068. 2022.
[44] Xunyi Zhao, Théotime Le Hellard, Lionel Eyraud- Dubbois, Julia Gusak, and Olivier Beaumont. Rockmate: an efficient, fast, automatic and generic tool for re-materialization in pytorch. In International Conference on Machine Learning, 2023.
13

---

A Checkmate
As previously introduced in the main section, Checkmate is a recomputation planner that operates using mixed integer linear programming. This appendix provides additional details on the internal workings of Checkmate.
Algorithm Let (N1,.... "Pl.....,P0) ,, where E N, be the original computation order. To define the i-th stage of computation (1 I S T), Checkmate uses a x T matrix { where R,, E {0. When given matrix R. the i-th stage of computation is subsequence of the original computation order, defined by (ny Rag = 1). The recomputation plan is then determined as the concatenation of stages. The maximum length the recomputation sequence is ( as a result of this formulation.
The objective of Checkmate to minimize -i= Ra,jo(n)), where c(n,) represents the cost of recomputing node nj. Additionally, Checkmate uses binary variables to determine whether to preserve the outputs of ny from stage i to i + and whether to free the outputs of n2 in stage i after evaluating T4g- The number of binary variables is O( + [W]M), where M is the number of in the DAG, Checkmate also introduces variables to represent the memory usage during the computation of node ny at each stage. Lastly, the memory budget constraint and the dependency constraints among the variables are added to formulate the recomputation problem as mixed- integer linear program (MILP).
A.1 Limitations Checkmate
Despite the ability of Checkmate consider node and value weights for any computation graph, it has several limitations.
Fixed topological order One such limitation is the requirement for fixed topological order. When a computational graph has multiple possible computation orders (i.c., the topological order is not unique), the memory usage can differ significantly based on the computation order selected. Identifying the topological order that minimizes memory usage is a difficult problem [33], and Checkmate may not be able to reduce memory usage significantly if the initial computation order is not memory efficient. As feasible recomputation sequences can become exponentially long for some graphs and budgets (27), it is challenging increase the number of stages to enable Checkmate to accommodate any feasible computation order.
LP-related limitations There are several additional limitations to Checkmate when the MILP relaxed to linear programming (L.P). Checkmate may fail to provide optimal solutions due to MILP/LP gaps, randomized rounding. and solver tolerances. In most cases, the LP relaxation method results in better objective values than the original MILP by allowing non-binary assign- ments for binary variables. However, Checkmate then applies randomized rounding variables that determine whether the outputs of ny at stage i is kept for stage í + 1, which can increase the objective value. The original paper suggests setting 0.9 x budget for the LP relaxed problem, considering the memory usage increase after randomized rounding. While often there is significant degradation in memory consumption, execution time may unexpectedly increase due to this process. We have observed that since the number of elements in { {Ra, } is O(IN)²), it is less likely have short recomputation sequences via randomized rounding when [A] is large. The degradation the recomputation cost is more likely to occur when the IP solver's tolerances are not small enough. However, determining the solver tolerances becomes a trade-off between convergence condition and optimization time.
the binary
B Application to extended problem settings
B.1
Computational graphs with node-intermediate memory usage
In this section, we explain some of the extended problem settings where our algorithm can be ap- plied.
14

---

Our algorithm can also be used find a memory- minimizing computation order without recom- putation. As mentioned in Appendix [A.]] the problem of finding a topological order of a DAG that minimizes peak memory is known to be NP-hard. We can use our algorithm to heuristically The original description in Section considers only the memory usage of input and output values. This setting reasonable for two reasons: first, deep learning frameworks and neural network com- pilers have shape propagation, making it easy to infer the sizes of intermediate tensors. Second, most cases, peak memory usage can be simulated with high precision even if node-intermediate memory usages are ignored.
However, precise peak memory estimation is necessary if the recomputation sequence must be highly optimized. For instance, if there is only a scarce amount of tensors alive when the mem- ory usage reaches its peak, we must consider node-intenee interintereeditee memory usage; otherwise, the es- timation error will be significant. Additionally, certain operators, such as matrix multiplication, temporal values may consume large amount of memory space. Considering node-intermediate memory usage is exceptionally important when dealing with grouped nodes. A grouped node may have substantially more vode-intermediate memory usage when it involves lot of computation. In our implementation, we estimate the node-intermediate memory of each grouped node from the temporal values produced and used within it.
Extensions in algorithm In the algorithm stated in Section 2.2 when a node n is added or re- moved from time t, we need to update the lifetime intervals for its input values and output values. To consider node intermediate values, we need an additional update: add (or subtract for removal) s' (n2) at time t, where s' (n) denotes the intermediate memory usage of n.
B.2 Memory minimizing topological ordering
solve this problem by prohibiting the addition and removal of nodes and allowing only node rotation
during mutations in the SA.
B.3 Offloading
Offloading, along with recomputation, is a popular method for reducing GPU memory usage. Of- floading involves moving data to secondary memory and back when it needed. Our algorithm is capable of optimizing offloading and recomputation simultaneously (although it could be applied to offloading alone).
The extended problem setting is as follows:
Given a computational graph and objective function f. the task is to construct a valid sequence .uV1...,,TT) and fetching sequence (Fi,...,Fr) FT), where F; C in;) that mini- mize the objective function. Here, F; denotes which input values are fetched (instead of us- ing previously produced ones) when computing ni. If all F3 are set to empty, the problem equivalent to the original recomputation. Note that the objective function f is now dependent on both S and the fetching sequence. For example, we can define the computational cost as Ei + Ever, c'(u) + Eve (n) c"(v), where e'(v) and c"(v) denote the cost of fetching and offloading, respectively. This objective function assumes that offloading and fetch- ing are done sequentially during computation. If background offloading and fetching are allowed, another suitable objective function may be defined.
Note that we only consider the timings of fetching in this problem setting, the optimal timings of offloading can be determined when they are given.
Extensions in algorithm We initialize all F; to be empty and allow an additional mutation in each SA iteration.
4. (Flip fetch) Choose random time t such that n, # nop and random value inputtn... If u E F,, update F; + F. Otherwise, update F, + F, U {v).
Unlike existing three mutations, this mutation is always valid, i.e., does not break the dependency of the sequence. Also, the inverse of this mutation is the mutation itself. The update of the life intervals
15

---

Graph
Op2
X
Lifetime
Op4
X
Remove recomputation
Add recomputation
Graph
Op2
X
Add
Op1
Add
Op3
Lifetime
Op4
X
Y
Figure 5: The update of lifetime by node addition or node removal. By adding a recomputation node, the life intervals of X and change. For either case, the update memory usage can be easily managed by the range-add range-max segment tree.
Figure 6: Example of node grouping. The blue node in the left graph fused together with its successors. The red nodes in the right graph are the fused nodes taking X as the input. Even without the SA, node grouping itself optimize recomputation since the last value will be always recomputed from X.
of can be done in the same way other mutations by considering the fetching and offloading as the addition or removal of special node which outputs from empty inputs.
C Technique for fast and better convergence
This section presents various techniques expedite the convergence of SA obtain better solutions.
C.1 Node grouping
node grouping a crucial pre-processing step for recomputing a sequence of nodes. Without node grouping, the SA has small chances of reaching solution that effectively reduces the cost the target function.
As mentioned in Section 2.2
Figure 6 provides example of node grouping. Even without the SA, node grouping can be used as simple recomputation strategy to lower peak memory usage. practice, our implementation uses the following algorithm. The do fu_(_s)() function in Algorithm determines whether to fuse
16

---

1e16
2.00
with grouping
1.75
without grouping
1.50
Objective value
1.25
1.00
0.75
0.50
0.25
0.00
0.0
0.5
1.0
1.5
2.0
Number of iterations
1e7
Figure 7: Memory reduction using node grouping for the LLaMA model. The blue line is the optimization with node grouping with the last 2M iterations being for decomposed nodes. The orange line is the SA without node grouping. The objective function is defined by equation (T) and the memory budget 6.7e+10, which is 10% of the original memory consumption.
node n, with its succeeding operators. We perform the fusion if the outputs of n, can be computed from inputs with smaller memory footprint, i.e., ZeEimpu((n)) s(v) < reontput(nn) s(v).
Algorithm Node grouping
Require: Valid computational sequence S (mu,...,,1) nT). Ensure: S is valid sequence with grouped nodes.
for .......77
. do
2: if do (o_fuuee(n)) then
3:
for j = 1..i+...... if output(n;)) S imput(n,) then n < grouped node of n; and n2.
4:
5:
6:
end if end for
8: Th nop
9: end
10:
end for
The node grouping algorithm with the do_ co_fusion function mentioned above has several charac- teristics. Firstly, Algorithm does not depend on the initial topological order of the computational graph and therefore works if the original sequence is memory-effiiient enough. Secondly, the above node grouping algorithm can optimally reduce the memory usage in specific simple cases, including the training graphs of sequential models, which were discussed in early studies of recom- putation techniques by Chen et al. [7]. In fact, the above node grouping algorithm can reduce the peak memory of such graphs
constant size providing that the all the values have the
size.
Figure 7 demonstratte the importance of node grouping in recomputation for LLaMA. When node grouping not used, the SA converged to final objective value 4.4 times larger than the result obtained by the SA with node grouping. On the other hand, using only node grouping (without any SA), were able to reduce memory usage for LLaMA to 7.8c+10, which is close to the 6.7c+10 budget, and lower than the optimized memory usage of 3.5e+1 achieved without node grouping.
C.2 LogAddExp segment tree
In our optimization approach, we assume that the objective function is function peak memory M(S) and computational cost C(S). For each mutation in the SA, the cost may be updated, but
17

---

1e14
Peak memory by Max reduction
4.5
Peak memory by LogAddEXp reduction
4.0
Objective value
3.5
2.5
2.0
1.5
0.0
0.5
1.0
1.5
2.0
Number of iterations
1e7
Figure 8: SA convergence for VIT model. We the first SA grouped nodes for 20M iterations and the second SA on decomposed nodes for 2M iterations.
the peak memory (S) changes only when the mutation affects the peak timings. Suppose there computational graph with memory peaks. In this case, MS) is determined by the higher of the two peaks, and mutations that reduce the memory usage around the second peak do not change (S). In general, multiple local memory peaks make it challenging to effectively reduce the objective value.
We can overcome this issue by employing sange-add/range-lggadeee range-add/range-max segment trees, introduced in Section 2 The range-add/range-maa segment tree recursively calculates max(a, b) for given interval and determines the maximum value within a range [4, r). On the other hand, the calculation of logaddexp(o, b) == log(exp(a) + exp(b)) instead of taking the maximum. The lo- gaddexp segment tree takes into account multiple memory peaks, enabling faster convergence of the SA.
segment trees instead of
segment tree uses the recursive
Ad Advantages and disadvantages While the logaddexp- segment tree expected to have bet- ter convergence when memory consumption is multimodal, it is essential to be wary of errors in memory usage simulation and numerical overflows. To avoid overflowing floating point operations, it is crucial to select appropriate base for log and exp. In the implementation used for Fig- ure g we re-scaled value sizes so that the maximum value is 10 and set the base to 2. However, this configuration may still cause overflow, particularly when [W| is large. There are important considerations to be made regarding the values held by the segment tree. In range-add/range- max segment tree, the t-the element of the segment tree represents the memory consumption time t. However, a logaddexp- based segment tree, this approach cannot be directly used since it leads to significant errors in calculating the peak memory. Specifically. unlike the max operator, logaddexp(x,,. -) log 2 + x >> I, meaning that repeatedly calculating logaddexp(a, b) for values and b close fact that the peak memory is only achieved for time where n, * nop. We initialized the elements of the segment tree with large negative constant and added large constant again when node is added at that time, ensuring more precise estimation of the peak memory.
will result in increasingly large values. In implementation,
utilized the
Figure shows the difference of the convergence by the kind of segment tree used in the SA state. The segment tree using logaddexp had slightly better convergence for the ViT model.
C.3 Implementation of segment tree
The role of the segment tree used in our algorithm is to efficiently calculate the peak memory, which can be obtained by the range-max query. Here, we introduce another implementation variant of the segment tree.
18

---

Dynamic segment tree The dynamic segment tree is a data structure that can efficiently handle the following operations: (1) Insert value index i. (2) Remove the value at index
(3) Add value z for range [1, r). (4) Calculate the maximum value for range [1. r). This kind of data structure can be created by extending binary search trees (for details, refer to texthooks such as (9). By using this segment tree, we do not need nop, introduced in Section 2.2 and can optimize computation sequence consisting of only N. However, the dynamic segment tree often involves heavy imple- mentation, the actual optimization time could increase. On the other hand, it is more efficient in terms of memory as it requires O(T)) space, where is the maximum length of the computational sequence during optimization.
D Experiments
In this section, additional information for the experiments done throughout the paper is provided to- gether with details on how to obtain computational graphs and the hyperparameters used to configure both, FastSA and Checkmate.
D.1 PyTorch Computational Graph
Optimizing PyTorch computational graph We extracted the computational graph using Py- Torch's symbolic tracing mechanism, which is available for PyTorch 2.0 or later. For our exper- iments, used PyTorch 2. 2....e220230404 as it had the necessary features for tracing unavailable in the latest stable release (v2.0.0).
Our recomputation algorithm was integrated by extending aot not_module. This module acts wrapper of model (torch. nn. Module) and optimizes the model's computational graph before the actual computation is done. Using symbolic tracing without real data, aot_module generates the computational graph for the model's forward and backward passes. then combines them into joint graph before optimizing and partitioning it for new model with optimized forward and backward computations. Users can pass a custom partition function (an optional argument) to aot_module take the optimization process.
We implemented custom partition function that follows these steps: receives the joint graph and applies the recomputation algorithm to create outputs from the same inputs. It returns the entire new computation graph as the forward module and configures the backward module to return only the pre-computed outputs assuming that all the tangents (i.e., gradient inputs) are one. The full graph is executed instead of separating into forward and backward passes as it is more memory-efficiem, When the graph is partitioned, some intermediate values for the backward pass need to be saved, and they remain in memory until the backward pass is complete.
computation graph that produces the same
Preprocess for Py PoTchh intermediate graph representation We perform Dead Code Elimina- tion (DCE) before applying the recomputation algorithms. Although our algorithm automatically does DCE during the SA by removing nodes, we perform explicit DCE in advance due to two rea- sons: dead code can affect node grouping performance and Checkmate does not support DCE.
Next, we convert the PoTch intermediate graph representation (torch. fx. (rrcc..xx.raph) to an internal graph format with explicit nodes N and values V. In torch. fx. Graph, the computational graph is represented as a list of nodes, and each value corresponds to a single node. If an operator has multiple output values in torch. fx..fx..raao Graph, it is represented as single node that outputs tuple, followed by get-item nodes that extract specific values from the tuple using an index. We can apply the recomputation algorithm to torch. fx. Graph intermediate representations, but it is easier to convert corch. fx. Graph to a more general graph frrmat to handle operators with multiple outputs efficiently. After the recomputation is done, convert format back to torch. fx. Graph.
For Checkmate, we follow the original paper's implementation and perform optimization directly on a torch. fx. Graph styled computational graph.
19

---

D.2
Checkmate settings
We re-implemented Checkmate 19 for the PyTorch intermediate representation. For Checkmate MILP, we used Gurobi as the internal MILP solver. For Checkmate LP, we employed PDLP solver, provided by OR-Tools [28] solve large-scale LP problems a multithreaded manner. PDLP was the fastest solver among the available options (CIP, GLOP, and SCIP) in OR-Tools.
The execution time and solution quality of Checkmate LP highly dependent on the hyperpa- rameters of the LP solver, as discussed in Appendix [A.]] We adopted the default values of le-4 for primal tolerance and dual tolerance in Checkmate's open-source implementation. We randomly generated 100 different thresholds in addition to the default threshold 0.5 for randomized rounding and selected the best solution the final output after testing these thresholds.
D.3 Environment
The proposed algorithm (FastSA) A) and Checkmate LP were evaluated using a cluster system, with each instance configured with CPU cores (Intel(R) Xeon(R) Platinum 8380 @ 2.30GHz) and 100 GiB of RAM with NVIDIA A10 80GB GPU. Although the PDLP solver fully uses the allocated resources, FastSA only requires a single CPU. Due to Gurobi license constraints, Checkmate MILP was executed on another machine with 36 CPU cores (Intel(R) Xeon(R) Gold 6154 CPU @ 3.00 GHz) and 376 GiB of RAM. The implementation of both FastSA and Checkmate was written in C++ and compiled using GCC 10.4.0 with the 03 option. The compiled module was integrated with PyTorch 2.130444 and the experiments were run using Python 3.9.12 with CUDA 11.8.
D.4 Settings of our algorithm
Initialization of SA state We initialized the default computation sequence S used in our SA as follows. First we fix the length of S as T = 2 and initialize all of them by nop. Then, for each n1 E of the default computation sequence (n....., n2),, we set ST4 +- ni. where / x
First-stage SA The SA algorithm in the first stage is applied to grouped nodes for efficient mem- ory usage reduction. The number of SA iterations is set to 20 million in this stage, but if the memory budget is met, the algorithm stops even the temperature is still high.
The initial temperature is set as 0.1% of the initial objective value, and the final temperature is set as 0.01% 1 of the initial temperature. The temperature the i-th iteration calculated as hę x exp(log( y (ha) x i/N), where h; and h denote the initial and final temperature, respectively, and N = 2 x 10? is the maximum number of iterations.
Second-stage SA The second stage of the SA algorithm is applied to decomposed nodes and targets to remove redundant computations. The number of iterations in this stage is 2 million. The initial temperature is set as the final temperature in the first-stage SA. The final temperature is set as 0.01% of the initial temperature, as in the first stage. The calculation of the temperature for each iteration is also the same.
D.5 Models
The models are obtained from timm (PyTorch Image Models) and transformers provided by Hugging Face. The versions of timm and transformers 0.9.1 and 4.28.1, respectively. All the models are set to the training mode. Random tensors are used for the inputs. For vision models shown in Table we set batchsize as 512. For GPU benchmarks, we reduced the batchsize 256 to compare the performance to the baseline, without recomputation. We summarize the all the models used in our experiments with their configurations below.
For text models shown in Table Al we used batchsize 128 and context length 512 to obtain the simulated memory and computational cost after applying recomputation. For the real GPU bench- marking, we reduced the values to 64 and 256, respectively.
For language models shown in Table 5 we used batchsize 8 and context length 2048, except for GPT2 with maximum context length 1024 to obtain the simulated values. For the real GPU bench- marking, we reduced the values to 4 and 1024 for all of them.
20

---

name
full model name
resolution
[
[V]
ConvNeXt [24]
convnext tiny convnextv2 large eva02 large_rre_____ cases__pachh444 224
224
932
1253
ConvNeXt V2 422
224
2840
3467
EVA-02 13 37
224
6205
6967
ViT 36
12
vit vir_large_ la__aarre__aachhl_____ 224
224
2817
3362
VIT (small)
vit_ walssmlllppphhht 2_44
224
1425
1706
MobileNetV3
[8
41
mobilenetv3 large 100
224
484
1155
EfficientNer
[38
efficientnet_hy aabaseepatchh__.. 224
224
996
1766
DeiT III
39
224
1545
1850
XCTT 1
xcit__in_11__p6__20
224
2818
3550
BEiT
beit beit_base si__base___aae_ faase__pachll..mm 224
224
1835
2187
CoAtNet [10
coatnet_2 2_rrs__________ __4
224
2204
3128
VGGI1 35
vggl1 resnet18
224
69
119
R 16
224
171
437
Table 3: List of timm models (PyTorch Image Models), [W] and [ denote the numbers of the
nodes and values the computational graph.
name
full model name
[W]
[V
ALBERT
22
albert-base-v2 bert-base-uncee uncased
2285
2419
BERT 11
2078
2386
DistilBERT 31
distilbert- bastttt-bb-euueee uncased -aseuus
1060
1228
ELECTRA
g
google/electra-smee small-ddisermminator
2099
2409
Table List of transformers text models.
name
full model name
|V]
GPT2 29
gpt2
1889
2179
GPT Neo 6
125M
EleutherAl/gpt-neo--P
125m
2088
2378
GPT Neo 2.7B
FleutterrAVggpt-eeo22277
5488
6238
BLOOM 32 560M
bigscience/bloomm56600
3667
4167
BLOOM 3B
bigscience/bloom-Shh
4567
5187
OPT 43 350M
facebook/optt350mm
3971
4581
OPT 6.7B
facebook/opt-6.1
5245
6059
LLaMA 40
7B
default config"
8330
8722
Table 5: List of large language models. *For LLaMA, we constructed model using default config- uration (7B parameters) provided by transformers amaConfig.
D.6 Solution stability of FastSA and Checkmate LP
Since both FastSA and Checkmate LP randomized algorithms, the recomputation plans depend on the seed. We executed FastSA and Checkmate I.P for various memory budgets and see how smoothly the solution change. Figure 9 shows the time-space tradeoff curve for different mem- ory budgets. Since Checkmate LP could not solve the instances with more than 2300 nodes, we compared the performance difference on relatively small models. For each model, set memory budgets as 0.15..0 0.20. 0.95. 1.00 and plotted the results. In most cases, FastSA found even better solutions than Checkmate LP. Although FastSA is randomized heuristic, its performance is stable and not heavily dependent random seeds. On the other hand, Checkmate LP's performance is not stable, mainly due to I.Prelated issues as discussed Appendix A. For instance, ViT (small), one of the solutions found by Checkmate LP had almost 10 times overhead. Although running Checkmate I.P with different seeds or hyperparameters may still yield good solutions, it requires tremendous computational resources.
21

---

ConvNeXt
VIT (small)
MobileNetV3
BERT
GPT Neo-125M
relative cost
Checkmate LP
1./5
1.50
1.25
1.00
10
1. relative memory
relative memory
relative memory
relative memory
relative memory
Figure 9: The time-space tradeoff curves for Checkmate LP and FastSA. For most settings, FastSA could find solutions with lower cost overhead than Checkmate LP for the same relative memory reduction. Checkmate LP is not stable due to LP relaxation. For MobileNerV3, FastSA could not find solutions for small memory budgets. This would be due to the limitation of the node grouping.
0.7
FastSA
0.6
Checkmate
Relative mermery
0.3
0.2
0.1
0.0
ConvNoXt
ComvNeXt V2
EVA-02
VIT
VIT (smail)
MobileNett/3
ElfccentNee
DelT III
XCIT
BEIT
CoAtNet
1.4
FastSA Checkmate LP
2.6
2.4
Relative
2.3
1.2
ContNeXt
ConvNeXt
EVA 02
VIT (smail)
MobiletietV/
EffcientNiet
DeiT
XCT
BEIT
CoANNet
space tradeoff for vision models with 25% (darker) and 50% (lighter) memory budgets. The batch size is 512. Checkmate LP was unable find solutions within 6-hour time limit for VIT and CoAtNet. The other models without Checkmate LP bars suffered from out-of- memory errors exceeding 100 GiB. The truncated results of the relative cost of Checkmate I.P for DeiT' and BEiT were 2.82 and 3.52, respectively.
Figure 10:
D.7 Full experimental results for comparison with Checkmate LP
The graph in Figure TO displays the recomputation results of our FastSA algorithm and Checkmate LP for all the vision models in Table 3 For VIT and CoAtNet, although Checkmate LP could run within 100 GiB RAM, it failed to find feasible solution within the 6-hour time limit. Further- more, the recomputation overhead by Checkmate LP for DeiT III and BEiT was more than 2.8 times, possibly due to randomized rounding of non-binary LP solutions. While FastSA obtained better solutions in both time and memory for ViT (small), DeiT III, and BEiT, it could not reduce memory much Checkmate LP for MobileNetV3 and EfficientNet, mainly due to suboptimal node grouping.
Figure 11 displays the recomputation results for all the text models in Table and recent language models in 3 For some text models, such as BERT and ELECTRA, our algorithm was unable to reduce memory usage sufficiently for a 25% budget, either due to limitations on the models
22

---

6.7
FastSA Checkmate
0.6
Relative memory
0.5
0.1
0.
0.2
ALBERT
DERT
DistuBERT
ELECTRA
GPT2
GPT Neo 125M
GPT Neo 2.78
BLOOM 560M
BLOOM 38
OPT 350M
OFT 6,78
LLaMA 78
1.8
FastSA Checkmate LP
2.6
2.5
2.4
Relative
1.2
0.9
ALBERT
BERT
DistiBERT
ELECTRA
GP12
GPT Neo 125M
GPT Neo 2.75
BLOOM 560M
BLOOM
OPT 350M
OPT 6.7B
LLaMA
Figure 11: Tpme- tradeoff for text models, including language models, with 25% (darker) and 50% (lighter) memory budgets. Checkmate LP failed for larger models due out-of aul-of-mmmory errors exceeding 100 GiB.
themselves or convergence to suboptimal solutions. However, overall, our algorithm was successful in finding good recomputation plans with only a small increase the recomputation overhead.
D.8 Actual GPU memory reduction and time increase
Throughout the above experiments, we discussed the memory decrease and cost increase by esti- mated values. This is because is hard assign exact computational costs for each node without execution (note that our computational graphs were traced without execution) and/or values may share the same physical storage actual GPU memory. In this section, we provide a summary of the results of actual memory usage and execution time measured and discuss the differences between the actual metrics and the simulated ones.
NVIDIA A A100 80GB GPU
Figure displays the difference of simulated and actual CUDA memory usages for various models for three cases: no optimization, 50% budget, and 25% budget. On average, the actual GPU mem- ory usages were less than the simulated values by 10% for no optimization cases, but could be to 5% more for 25% budget cases. most cases, the simulated memory usages were higher than the allocated usages, mainly because the memory simulation is not precise enough. Notably, operations like reshape, expand, and slicing may produce new views of the original tensor and do allocate new GPU memory. Our memory simulation does not consider this information, leading to overesti- mation the memory usage. Memory aliasing, tensor views, and node-intermediate memory usage can account for these differences. For the 25% budgets, the simulated and allocated memory usages are similar for all models. However, in some instances, the allocated memory exceeds the simulated memory, which may be due to inaccurate memory simulation. One of the reasons is that do not consider the node-intermediate memory usage for simulation, which underestimates the mem- ory usage. We can address this issue in problem setting by simulating this value (please refer to Appendix B.I.
Figure 13 shows the difference between the actual increases in execution time and the simulated ones by the unit cost. We observed that the simulated cost is relatively close to the actual measured overhead. The difference in simulated time increase from the actual time increase was at most 7% and 12% for 50% and 25% memory budgets respectively. We observe that practical recomputation plans can be successfully obtained for many models using the unit cost. However, the reason behind
23

---

simulated
1.0
Relative memory
0.8
0.6
0.4
0.2
0.0
viT
MobileNetV3 Efficient
III XCIT
BeiT
ALBERT
BERT
DistilBERT
ELECTRA
GPT2
GPT Neo 125M BLOOM 500M
350M
The actual memory usage of optimized models by FastSA. For each model, the three bars indicate the actual memory usage of the model when executed without optimization, and with 50% and 25% budget constraints. The simulated memory is indicated by the red line. The actual memory is the total size of CUDA tensor allocated by the PyTorch CUDA allocator. As detailed in Section D.5 the batchsizes and context lengths are modified to run the original model without recomputation for comparison.
Figure 12:
Relative computation time
1.5
simulated
1.4
actual
1.2
1.1
1.0
0.9
ViT
MobileNetV3 Efficiettiggt XcIT
III
BelT
ALBERT
BERT
DistilBERT
ELECTRA GPT2
GPT Neo 125M
BLOOM OPT 350M
Overhead of execution time by recomputation. For each model, the two bars correspond to 50%, and 25% memory budget setting, respectively. The simulated increase of the computational cost indicated by the blue line.
Figure 13:
similar simulated and actual overhead is balanced mix of operators with high costs, such as matrix multiplication, and operators with negligible costs, such reshape, effectively averaging the cost. By accurately profiling the operators, it is possible to obtain more efficient recomputation plans.
24


Reviews:
[{"review_id": "ZzLzCCzpLp", "reviewer": "Reviewer_w4oG", "summary": "This paper introduces a fast simulating annealing combined heuristics approach for gradient checkpoint/recomputation. \nThe solution achieves a significant memory reduction of 73% with an average recomputation cost of 18%. \nIt outperforms the state-of-the-art MILP-based technique Checkmate in terms of runtime by orders of magnitude, while still meeting the memory budget requirements.\n\nUnlike Checkmate, which utilizes MILP optimization on the graph, this work employs a sequence of nodes to represent values and a binary variable LS(v, t) -> {0, 1} to indicate if a value v is in memory at time t. The optimization objective is defined as the maximum of the resource utilization function M and the budgeted compute cost C. To optimize the solution, the paper utilizes an add-max segment tree structure for efficient interval evaluation. The optimization algorithm employed is simulating annealing (SA). The approach allows for three mutations in the sequence: add, remove, and rotate, enabling effective exploration of different configurations.", "questions": "1. It would be important to re-evaluate the Checkmate results using Gurobi. \n\n2. Can you show a training time comparison of the solutions generated by SA and Checkmate? \n\n3. Can you provide some explanation of why Checkmate as an ILP failed to meet the memory budget requirements? Are the constraints not precisely formulated in Checkmate?\n", "limitations": "- The grouping heuristics for nodes might not be optimal\n- There are many other GPU memory usage factors to consider in the optimization. ", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The approach presented in this work introduces a novel and scalable solution for addressing the recomputation problem. Rather than formulating the problem as MILP which can be computationally expensive, this work reformulates it using sequences of nodes and utilizes an efficient data structure to evaluate its objective function. The problem is then optimized using the well-known simulated annealing algorithm, eliminating the need for a commercial solver.\n\n2. By adopting this approach, this work manages to reduce the memory requirements to meet the specified 25% and 50% memory budget without incurring high computation costs.  \n\n3. The optimization time of this approach is notably faster, around 3-4 orders of magnitude compared to the previous state-of-the-art technique Checkmate. \n", "weaknesses": "1. The baseline Checkmate is run using the open-source OR-tool instead of Gurobi. It not only affects the runtime (which can be orders of magnitude longer) but also the quality of the results of MILP. \n\n2. The comparison of end-to-end training time between the proposed approach and Checkmate is not provided. The increase in compute cost may not necessarily lead to improved performance if the workload is not compute-bound. \n\n3. The topic could probably be better evaluated and discussed at a systems conference. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. In our original experiments, we compared our algorithm to Checkmate LP, instead of the exact ILP, using open-source LP solver OR-Tools due to the prohibitive cost of commercial ILP/LP solvers. However, we acknowledge that the evaluation of the solution quality of our algorithm is important and we procured a trial license of Gurobi and conducted an initial exploratory study. Due to the short rebuttal period, we did not replace all of the Checkmate LP results solved by an OSS solver with Gurobi. Instead, we concentrated on verifying the optimality of FastSA’s solutions by solving Checkmate ILP (no relaxation) with Gurobi. We will replace the results in the next version of the manuscript.\n\nAlso as summarized in the global author rebuttal, we first tried to evaluate all of the models used in our original experiments by Checkmate ILP with Gurobi, but even for smaller models within 1000 nodes, none of them was solved within the time limit of 3600 seconds. Therefore, we prepared much smaller instances (vgg11, resnet18) for calculating optimal solutions.\n\nTable 2 in the attached pdf shows the optimality gap between FastSA and Checkmate ILP. For vgg11, FastSA’s recomputation plan contained at most one more node than Checkmate optimal. For resnet18, FastSA’s plan contained at most 4 more nodes than Checkmate optional, where the original computation graph has 171 nodes. In addition, Table 1 summarizes the experiment results by adding FastSA's results to the table in the paper of Moccasin, a recently published recomputation algorithm at ICML 2023. Here, we can also discuss the solution quality of FastSA, compared to Checkmate/Moccasin optimals. In the case of RL (random layered), FastSA's results are better than Checkmate's or Moccasin's. This is because there is a high degree of freedom in model topological order in the RL cases, and FastSA can simultaneously optimize the topological order while minimizing the memory footprint with almost no additional recomputation nodes.\n\n**Q: It would be important to re-evaluate the Checkmate results using Gurobi.**\n\nA: Please kindly refer to the above comments and the general response for more details on this.\n\n**Q: Can you show a training time comparison of the solutions generated by SA and Checkmate?**\n\nA: In our paper, we did not directly assess the end-to-end training time. Instead, akin to prior work on recomputation [Kumar et al. NeurIPS 2019, Kusumoto et al. NeurIPS 2019, Jain et al 2020, Bartan et al. 2023], we focused on quantifying the additional computational time required when the memory budget for the model is constrained. However, an estimation of end-to-end training time can be inferred from our optimization results. Generally, the memory used in training directly corresponds with the batch size. Hence, if the memory is reduced to 1/k, it enables handling a batch size that is roughly k times larger. Assuming an average increase in computational time per operator, symbolized as c (a value estimable through benchmarks), the throughput of end-to-end training is likely to improve by a factor of k/c.\n\nFor distributed training, data transfer bandwidth can pose a bottleneck rather than the arithmetic computational cost, rendering the end-to-end experimental setup more complex. However, as detailed in Appendix B.3, FastSA can be adapted to problem settings that include data transfer. With a suitably defined objective function, it may potentially accommodate these scenarios in the future.\n\n**Q: Can you provide some explanation of why Checkmate as an ILP failed to meet the memory budget requirements? Are the constraints not precisely formulated in Checkmate?**\n\nA: The experiment results of Checkmate in our paper are all solutions from the LP relaxation problem setting, not the ILP. Since SCIP, an open-source ILP solver, could not solve any instance in our original experiments within the time limit of 6 hours, we decided to solve the relaxed problem. The solution obtained from LP relaxation is generally not an integer, and the memory budget constraint may be violated during the randomized rounding process. More detailed explanations about this can be found in Section 5.2 of the original paper of Checkmate [Jain et al. 2020].\n\nAs detailed previously, to evaluate the solution quality compared to the ILP optimal, we added extra results on vgg11 and resnet18 in the attached pdf file. Please note that even with Gurobi, we could not find optimal solutions within 3600s even for the smallest models used in our original experiments.\n\nWe hope this adequately addresses your concerns and appreciate your feedback."}, {"review_id": "di5oHOb45p", "reviewer": "Reviewer_kH45", "summary": "This paper proposes the Fast Simulated Annealing Algorithm (FastSA), based on the Add-max segment tree and simulated annealing, to optimize memory usage and training time. Furthermore, FastSA introduces grouped nodes to aid the convergence of simulated annealing and effectively reduce the peak memory. It can also be extended to incorporate other memory optimization techniques such as offloading. As a result, FastSA successfully finds optimal solutions, even for large language models where existing checkmate algorithms fail, leading to optimized training speed and memory usage.", "questions": "If a proper comparison between this algorithm and Moccasin is conducted and it is demonstrated that FastSA achieves significantly better performance, I am willing to raise the score to 7 or higher.\n\n+ A qualitative comparison between the Moccasin algorithm and FastSA, highlighting the superior aspects of FastSA.\n\n+ A quantitative comparison between the Moccasin algorithm and FastSA based on experimental results.\n\n", "limitations": "The authors provide a thorough discussion of the limitations and potential future work in the paper.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper presents a **novel approach** that differs from the existing checkmate method in three key aspects. \n+ First, it formulates the problem as optimizing the sequence of nodes by defining each node to represent a single operator and introducing the concept of Lifetime. \n+ Second, it proposes an algorithm that effectively combines the add-max segment tree and simulated annealing to solve this problem. \n+ Third, it optimizes the process through grouping. \n\nThese ideas are not only highly innovative but also demonstrated to **successfully optimize memory and training time when applied to large models**. This highlights the significance of the findings in this paper.\n", "weaknesses": "**Lack of comparisions**\n\nThe FastSA finds the optimal solution based on the concept of lifetime. This approach shares many similarities with the recently proposed event-based optimization algorithm called Moccasin [1]. Specifically, Moccasin also claims to achieve superior performance compared to MLIP-based Checkmate by utilizing time interval information for optimization. However, the FastSA paper does not mention this and lacks quantitative or qualitative comparisons.\n\n[1] Moccasin: Efficient Tensor Rematerialization for Neural Networks", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\nWe acknowledge and appreciate the reference to Moccasin, which was unavailable at the time of our paper's submission, but is indeed an important comparison to make. Consequently, we have conducted further comparisons between Moccasin and FastSA, which we will include, along with a detailed explanation, in our revised manuscript.\n\n**Q: A qualitative comparison between the Moccasin algorithm and FastSA, highlighting the superior aspects of FastSA.**\n\nA: Different from Checkmate and FastSA, Moccasin uses Constraint Programming (CP) to address the rematerialization problem. By focusing on lifetime intervals, Moccasin reduced the number of integer variables in CP to linear size to the number of nodes, compared to the quadratic number of boolean variables in Checkmate. However, the recomputation plans produced by Moccasin may not be efficient enough, because its search space for the solutions is $O(2^{Cn} + n^{Cn})$ and is significantly narrower than Checkmate’s $O(2^{n^2 + nm})$ or FastSA’s, where C is the hyperparameter representing the maximum number of recomputations for each node and in the open-source implementation, it is configured to C=2. The solution space of FastSA is even larger than Checkmate, meaning that FastSA can find even better recomputation plans than both Checkmate or Moccasin. Setting C=2 is practical in terms of execution time, but then the memory reduction on sequential models is limited to O(√n), whereas it is feasible to achieve O(log n) or O(1) memory by allowing more than two recomputations for each node.\n\nIn fact, the experiments in the original paper of Moccasin (also cited in Table 1 of our attached pdf file), were conducted for 90% and 80% memory budgets, where the peak memory can be reduced with a small number of recomputations. However, as shown in our original experiments, we can even reduce memory to 25% given modern large-scale models. In such situations, the constraint of C=2 in Moccasin can significantly limit the solution quality. By setting a larger C value, the space of candidate solutions increases exponentially, limiting the potential applicability of aggressive memory reductions. For more explanation on the size of search space, please refer to Table 1 of Moccasin paper.\n\n**Q: A quantitative comparison between the Moccasin algorithm and FastSA based on experimental results.**\n\nA: As shown in Table 1 of our attached pdf, FastSA consistently identified solutions with a smaller Total Duration Increase (TDI), which refers to the recomputation overhead, faster than Moccasin in many instances. Notably, for the random layered (RL) case, FastSA successfully minimized memory consumption up to 4.9% smaller TDI than Moccasin without the need for additional recomputation nodes, simply through optimizing the topological ordering. While Moccasin operates faster than Checkmate, it takes approximately an hour for cases nearing 1000 nodes, whereas FastSA can find superior solutions in a fraction of this time.\n\nIt's important to note that constraint optimization problems, as addressed by Moccasin and Checkmate, generally demand a substantial computational time to arrive at feasible solutions, more so when the solution space shrinks. As a result, even with equal node numbers, the computational time escalates considerably as the budget tightens. This is evident in Table 2 of our additional materials, which illustrates an example of Checkmate results. Checkmate took more than 40x longer optimization time for resnet18 when the budget was tightened from 80% to 60%. Accordingly, the benefits of employing FastSA are even more pronounced when navigating tighter budgets as the optimization time only increased up to 5% when the budget was changed from 90% to 50% in the case of resnet18.\n\nWe trust that we have adequately addressed your concerns. Thank you once more for your thoughtful review."}, {"review_id": "URp33JaYbV", "reviewer": "Reviewer_a8KW", "summary": "This paper proposed a new method for optimizing recomputation in neural network training. It formalizes the recomputation as a sequence of nodes that indicate the computation schedule. It then uses simulated annealing with segment tree to find a sequence that optimizes throughput with a certain memory budget. Experiments show significantly reduced computation overhead compared to the state-of-the-art optimizer checkmate in memory restricted cases. It also has a significant lower solving time than checkmate.", "questions": "1. Figure 3: Why is there no checkmate bar for AVE?\n2. Figure 3: Why is there not much difference between FastSA and Checkmate for 50%?\n3. Is it true that FastSA cannot beat Checkmate on transformers?", "limitations": "The linearization of the graph and the relaxation used in the algorithm makes the solution sub-optimal.\n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is clean and easy to follow.\n2. The formulation of the problem is clear.\n3. The proposed method is novel for this problem.\n4. The evaluation is instructive and convincing. The real implementation is aligned well with the simulator.", "weaknesses": "1. The contribution is incremental. It does not open a new problem or new angle, and the technique is common.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\nAcknowledging that our problem formulation shares resemblances with existing work, it's pivotal to highlight that, based on our knowledge, our algorithm is the first algorithm able to perform recomputation in arbitrary computational graphs and can be applied to real-world models involving thousands of nodes making our solution truly innovative and unique in this area. Existing algorithms such as Checkmate were capable of handling prevalent models of their time, such as ResNet. However, with the evolution of model sizes and complexities, these methods fall short when it comes to handling larger models like modern Transformers, which are heavily constrained by the memory capacity of contemporary devices. So, in this sense, our research tackles a problem that's becoming increasingly significant, offering new solutions and greatly pushing the current boundaries of what's achievable.\n\nAlthough simulated annealing (SA) would be a common idea to solve optimization problems, we utilized segment trees within our SA to reduce the time complexity of re-evaluating the peak memory after the mutations on the computation graph from naive O(n) to O(log n). This technique is the core of our algorithm allowing it to process large graphs over thousands of nodes.\n\n**Q: Figure 3: Why is there no checkmate bar for AVE?**\n\nA: The absence of an AVE bar for Checkmate in the figures is due to Checkmate's inability to find a solution within the six-hour time limit for some instances. This detail is covered in line 252 of the manuscript, but we acknowledge that the explanation might not have been sufficiently clear. For better clarity, a direct explanation will be added in the figure captions.\n\n**Q: Figure 3: Why is there not much difference between FastSA and Checkmate for 50%?**\n\nA: The relative ease of recomputation at this level allows both algorithms to effectively identify good solutions. However, as the budget reduces to 25%, the number of feasible computation orders decreases and the recomputation plans will become more complex, introducing challenges for rematerialization algorithms. This point will be better elaborated in the revised manuscript for improved understanding.\n\n**Q: Is it true that FastSA cannot beat Checkmate on transformers?**\n\nA: When dealing with transformer models, FastSA tends to outperform Checkmate. This is supported by a comparative analysis in Figure 12 of the supplementary materials for transformer models. For the six transformer models in Figure 12 where both FastSA and Checkmate results are available, FastSA has on average 9.6% smaller memory and 4.1% faster than Checkmate."}, {"review_id": "7EjNPbPlFc", "reviewer": "Reviewer_wXFY", "summary": "This paper proposed an algorithm for the rematerialization problem. The proposed algorithm is based on simulated annealing.", "questions": "- As far as I know, Checkmate has two versions: 1) MILP, 2) Rounded LP. It looks like that the authors consider the rounded LP only. Based on my experience, rounding leads to very low quality solutions especially for complex DAGs. I understand because of gurobi being a commercial solver, it could be difficult for the authors to obtain a comparison for the gurobi version. One solution to this is to use the SCIP solver (maybe use directly its API or via CVXPY) or even use the available MILP solver in Google Or-tools. More importantly, what I want to get at is that it's not clear to me how close to the optimal the solution produced by the proposed algorithm is. This could be numerically studied by taking a smaller graph (maybe 100 nodes or so) and then solving the MILP of Checkmate using SCIP to global optimality. Then, the authors could report how close to this number their solution is. I expect this would make the results more complete.\n\nIn summary, I am convinced that the proposed method works better than rounded LP, but also want to point out that rounded LP doesn't work well anyways. I'd be curious to see how far it's from optimal at least for graphs of manageable size.\n\n- Line 202: \"Since the exact cost of a node cannot be determined before execution, we assumed that all operators have the same cost\". Perhaps the authors consider profiling the operator durations? Or could it be possible to come up with a more sophisticated (than unit cost) way of assigning durations to the operators?\n\n- I am not expecting the authors to be familiar with this paper or include it in their numerical comparisons since it came out very recently, but it's quite relevant: [r1] \"Moccasin: Efficient Tensor Rematerialization for Neural Networks\" from ICML 2023. The authors of [r1] show that there exists an nlogn-variable formulation (as opposed to the n^2-variable formulation of Checkmate).\n\n- I think this claim in line 311 may need to be made more precise: \"This complexity scales quadratically with the size of the input computational graph, making it a time and memory consuming process.\" The number of decision variables scales quadratically, not the overall algorithm complexity. The overall algorithm complexity is much worse than quadratic since it's a MILP.\n\n- input(n), output(n): Maybe a more standard terminology is to use parents/children predecessors/successors?\n\n- I apologize if this is already mentioned in the text but, what is the difference between \"FastSA\" and \"FastSA Only\" in Figure 4?\n\n- How do you define the optimization time in Figure 4? This looks a bit confusing to me, since we know that by running a milp solver with infinite time will get us optimal. ", "limitations": "Please see the questions section.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Please see the questions section.", "weaknesses": "Please see the questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We really appreciate your feedback emphasizing the importance of comparing FastSA results to Checkmate ILP. In our original paper, we relied solely on Checkmate LP due to licensing issues and the large sizes of the models. However, we now understand the necessity of a comparison with Checkmate ILP to evaluate FastSA's solution quality more effectively.\n\nAs detailed in the global author rebuttal, we obtained a trial license of Gurobi and carried out additional experiments. Given Checkmate ILP's inability to find feasible solutions for the models used in our original experiments within a 3600-second time frame, we resorted to using vgg11 and resnet18. As per your suggestion, these models, which are composed of approximately ~100 nodes, were used to compute optimal recomputation plans.\n\nTable 2 in the attached PDF illustrates the optimality gap between FastSA and Checkmate ILP. For vgg11, FastSA's recomputation plan at most includes one additional node than the optimal plan identified by Checkmate. For resnet18, FastSA's plan at most contains 4 more nodes than the optimal Checkmate plan, in a computation graph originally consisting of 171 nodes.\n\nAdditionally, Table 1 summarizes the experimental results, incorporating FastSA's results alongside those from the Moccasin paper. The quality of FastSA's solutions can be evaluated here in comparison with the optimal solutions derived by Checkmate or Moccasin. FastSA outperforms both in the case of random layered (RL) models, owing to the substantial flexibility offered by these models in terms of their topological order. FastSA can concurrently optimize this topological order whilst minimizing memory footprint, with almost no additional recomputation nodes.\n\nAlthough we had to significantly reduce the model size as depicted in Table 2 to obtain the exact solution through MILP with Checkmate, we anticipate that FastSA presents promising opportunities for producing superior solutions than Checkmate for complex and larger models due to the greater degree of flexibility in the model's topological order.\n\n**Q: Line 202: \"Since the exact cost of a node cannot be determined before execution, we assumed that all operators have the same cost\". Perhaps the authors consider profiling the operator durations? Or could it be possible to come up with a more sophisticated (than unit cost) way of assigning durations to the operators?**\n\nA: Indeed, by utilizing profiled runtime or flops estimation as the node cost, a more precise recomputation plan can be derived. However, there is a limitation when naively profiling a computation graph — the initial memory footprint of that graph should fit within the GPU memory. This becomes a challenge when recomputation is indispensable to fit the model within a single GPU. To negotiate such scenarios, an initial run of FastSA with a unit cost can be performed to decrease the memory usage, followed by profiling, and after that a more suitable recomputation plan can then be developed based on the profiled runtime.\n\n**Q: I am not expecting the authors to be familiar with this paper or include it in their numerical comparisons since it came out very recently, but it's quite relevant: [r1] \"Moccasin: Efficient Tensor Rematerialization for Neural Networks\" from ICML 2023. The authors of [r1] show that there exists an nlogn-variable formulation (as opposed to the n^2-variable formulation of Checkmate).**\n\nA: Please check the above reply and the global author rebuttal.\n\n**Q: I think this claim in line 311 may need to be made more precise: \"This complexity scales quadratically with the size of the input computational graph, making it a time and memory consuming process.\" The number of decision variables scales quadratically, not the overall algorithm complexity. The overall algorithm complexity is much worse than quadratic since it's a MILP.**\n\nA: We apologize for the unclear explanation. We will revise it in our final version.\n\n**Q: input(n), output(n): Maybe a more standard terminology is to use parents/children predecessors/successors?**\n\nA: We will revise that section in our final version.\n\n**Q: I apologize if this is already mentioned in the text but, what is the difference between \"FastSA\" and \"FastSA Only\" in Figure 4?**\n\nA: “FastSA Only” means only FastSA could find a solution and Checkmate LP could not find it within the time limit of 6 hours. We will add the explanation to the caption.\n\n**Q: How do you define the optimization time in Figure 4? This looks a bit confusing to me, since we know that by running a milp solver with infinite time will get us optimal.**\n\nA: Optimization time in Figure 4 is not the total node costs of the optimized computation graph, but the time spent by FastSA/Checkmate algorithm to find a suitable recomputation plan. \n"}]

Author rebuttal:
Dear reviewers

We appreciate the insightful feedback. The review comments have been proven to be very valuable and help to increase the quality of the paper. Here, we provide additional results regarding two common concerns; (1) Comparison with Moccasin [Bartan et al.], a rematerialization algorithm recently published in ICML 2023, and (2) The solution quality of FastSA (our algorithm), compared to Checkmate when solved with the exact ILP (integer linear programming), using a strong commercial solver, Gurobi.
For other questions, please refer to the individual replies.

**Q: Comparison of FastSA and Moccasin [Bartan et al. ICML 2023]**

A: We have benchmarked FastSA against Moccasin using the publicly available data that the authors released and observed that FastSA is ~100 times faster and could find solutions with less or equal recomputation overhead than Moccasin for 10/12 cases when evaluated in graphs with ~1000 nodes.

Moccasin formulates the problem of rematerialization by constraint programming (CP), which is similar to Checkmate’s ILP. Compared to Checkmate, Moccasin reduces the number of boolean variables from quadratic to linear, resulting in faster execution. 
However, the scope of feasible solutions within Moccasin is significantly limited owing to the role played by the C hyperparameter, which represents the maximum number of times a value can be rematerialized. In practice, C is set to 2 because the search space of Moccasin is $O(2^{Cn} + n^{Cn})$ and thus its computational time grows exponentially for larger C. Specifically when dealing with sequential models, this setting of C=2 imposes a limit on the memory reduction to O(√n), despite the existence of recomputation sequences with O(log n) or O(1) memory utilization that can be achieved by both FastSA and Checkmate.

Table 1 in the attached pdf file shows the comparison among Checkmate, Moccasin, and FastSA. For RL (random layered) graphs, FastSA is able to find better solutions than both Checkmate and Moccasin. Especially, for 90% budget cases such as RL1 and RL2, FastSA reduced the memory budget up to 22.6% just by optimizing topological ordering without adding extra recomputation nodes.

**Q: How close to the optimal the solution produced by FastSA is?**

A: In the original version of the FastSA paper, a comparison with Checkmate ILP was not performed due to the prohibitive cost of commercial ILP solver licenses. Additionally, open-source solvers such as SCIP failed to complete any of the experiments within a reasonable time frame. However, as pointed out by reviewers wXFY and w4oG, a rigorous comparison with an ILP solver is necessary to ensure fairness.
To accommodate this, we procured a trial Gurobi license and reiterated all experiments on relatively small models (within 1000 nodes, 50% budget) in the evaluation section of the paper after cross-verifying our Checkmate ILP+Gurobi implementation with the Checkmate paper experiments. Regrettably, even with Gurobi, the experiments could not be completed within a time limit of 3600 seconds due to the considerable size of the models involved. As per reviewer wXFY's suggestion, we have incorporated graphs with approximately ~100 nodes such as vgg11 or resnet18 and obtained their optimal solutions using Checkmate ILP+Gurobi.

Table 2 in the attached PDF presents a comparison between Checkmate and FastSA for smaller instances. For vgg11, under the tightest budget constraint (80%), FastSA identified the same recomputation plan as Checkmate. For resnet18, wherein the tightest budget that allowed an optimal solution was 60%, FastSA's Total Duration Increase (TDI), which means recomputation overhead, was 1.2% above Checkmate's TDI. Even though FastSA did not determine an optimal recomputation plan for this scenario, its runtime is almost 100 times faster than Checkmate's, suggesting its potential for integration into neural network compilers to optimize a diverse range of large models

Table 1 also compares Checkmate ILP and FastSA. For the RL cases, FastSA found solutions with a maximum TDI of 0.3%, whereas all Checkmate solutions had a maximum TDI of 0.8%. In these scenarios, FastSA's solutions superseded those of Checkmate as FastSA was successful in reducing the peak memory through optimizing the topological ordering. Please refer to Appendix A.1 for a more detailed explanation of this constraint in Checkmate when dealing with computation graphs that offer high topological flexibility.

Meta-review: The paper proposes a new method to optimize recomputation during neural network training (recomputation = a sequence of nodes in some computation graph that indicate the computation schedule during training). It's based on simulated annealing strategies that attempt to find a sequence that optimizes throughput with a certain memory budget. The paper boasts memory reductions by 73% with mild (18%) computational overheads on average.

The reviews were largely positive. The main concerns raised were that: 1) similar techniques have appeared in the related literature; 2) the topic might be more appropriate for a MLSys style venue.

On balance, the method shows enough improvements, and there's enough interest in NeurIPS in training big models, that including it in the program makes sense.
Decision: Accept (poster)