SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For the  $50\%$  memory budgets, FastSA algorithm was successful in meeting the target budget for all models, with an average increase in model execution time of  $7\%$ . However, for the  $25\%$  budgets, FastSA could not reduce memory to the budget for some models due to either limitations on the model itself or convergence to a suboptimal solution due to SA heuristic limitations. Despite this, 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 Checkmate LP in reducing memory usage, particularly for larger neural networks.
A comparison with Moccasin was also conducted by applying FastSA to the publicly available data in [3]. Table1 presents the results from the original Moccasin study, extended to include outcomes for FastSA. The experiments pertain to two graph variants, namely random-layered (RL) and graphs mentioned in [19] (CM). For the RL graphs, FastSA achieved the lowest recomputation overhead
3The PDLP solver is used for Checkmate LP due to the licensing costs associated to Gurobi. Checkmate LP experiments were executed in parallel using a cluster environment. We verified that for the experiments in Figure 4, the PDLP solver could solve all the instances solved by Gurobi, making it a reasonable open-source alternative.
Graph	(n,m)	B	CHECKMATE MILP	MOCCASIN	FASTSA
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	no experiment	no 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	no 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	3.4	80.0	3611.8	2.5	78.4	36.3
70	no experiment	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	no experiment	no experiment	0.8	66.8	11.0
in all instances. Remarkably, for RL1 and RL2 cases, FastSA managed to decrease memory usage without adding new recomputation nodes, through the optimization of the execution's topological ordering.
3.3 Solution optimality
Table 1: Comparison with Moccasin. The table includes results of Checkmate MILP and Moccasin from [3], extended with FastSA data. The columns B and CI represent memory budget and cost increase percentage, respectively. Alongside original results for  $90\%$  and  $80\%$  budgets, a  $70\%$  budget row only demonstrating FastSA results is also inserted. For all random-layered (RL) cases, FastSA exhibited the smallest CI, as it optimizes topological ordering, reducing memory without materialization.
Model	(n,m)	B	CHECKMATE MILP	FASTSA
CI	Mem	Time	CI	Mem	Time
VGG11	(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 on small models. For the resnet18 case with  $50\%$  budget, the best feasible solution was used because the MILP solver could not find the optimal solution within the time limit.
We conducted a study on how closely FastSA solutions match with the optimal solution by comparing 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 a 50% budget, FastSA managed to cut down memory use under the budget in all cases, but with up to 2.4% more overhead from recomputation compared to Checkmate MILP. Even though FastSA fell short of finding the optimal solutions for these cases, it was up to 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 Appendix A.1). This can be found in the results for RL graphs in Table 1, where FastSA could find better solutions than Checkmate MILP or Moccasin.
4.1 Model Parallelism
One common approach to scaling a 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 overhead required for each iteration of the model. Rajbhandari et al. [30] found that these approaches can achieve only  $5\%$  of the performance peak of a V100 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 a 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. [7] as a 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 a mixed integer linear programming (MILP) problem and proposed Checkmate, a solver to find an optimal recomputation plan. Checkmate tries to minimize 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 a hyperparameter 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 FastSA. Also, in a parallel 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 a smaller graph than the entire model, and then finds the recomputation plan across the layers by Rotor [17], 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 a single CPU core, taking less than four seconds on average. This significant reduction in computational time is achieved by leveraging efficient heuristics and optimization techniques. Our results demonstrate that our 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 required. Beaumont et al. [4] discuss the combination of offloading and recomputation as a means of further reducing memory utilization. Our proposed method can easily be extended to support
offloading, as detailed in appendix B.3. Another approach that leverages offloading is the ZeRO technique proposed by Rajbhandari et 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 a 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 computations and weight storage up to 2-bits in 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 a memory-efficient architecture that can perform backward computation without saving the activations. However, its applicability is limited to residual networks only. It is worth noting that these techniques are largely orthogonal to our proposed method and can be combined to further improve memory savings in neural network optimization.
In this paper, we present a novel method for recomputation that offers several key advantages over existing approaches. Our method is applicable to general graphs and can 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 a 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 in a highly efficient manner. This makes it ideal for integration into neural network compilers, where it can further streamline the optimization process. Additionally, our algorithm is highly flexible and can handle a wide range of problem settings, including recomputation with offloading and consideration of node-intermediate memory usage.
We explored the tradeoff between computation time and memory usage and evaluated the effectiveness 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 a 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 to 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 computations. 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 a gap between the simulated and actual GPU memory usage as shown in Appendix D.8.
Overall, our proposed approach offers a fast and efficient method for neural network optimization, providing significant improvements over existing techniques. However, future research may be required to support a wider variety of computation.
[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.
[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, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, et al. Meshtensorflow: 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.1556, 2014.
[36] Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint arXiv:2106.10270, 2021.
[37] Quan Sun, Yuxin Fang, Ledell Wu, Xinlong Wang, and Yue Cao. Eva-clip: Improved training techniques for clip at scale. arXiv preprint arXiv:2303.15389, 2023.
[38] Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In International conference on machine learning, pages 6105-6114. PMLR, 2019.
[39] Hugo Touvron, Matthieu 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, Timothee Lacroix, Baptiste Roziere, 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 arXiv:2110.00476, 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 arXiv:2301.00808, 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 transformer language models. arXiv preprint arXiv:2205.01068, 2022.
[44] Xunyi Zhao, Théotime Le Hellard, Lionel Eyraud-Dubois, 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.
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  $(n_{1},\ldots ,n_{T})$ , where  $n\in \mathcal{N}$ , be the original computation order. To define the  $i$ -th stage of computation ( $1\leq i\leq T$ ), Checkmate uses a  $T\times T$  matrix  $\{R_{i,j}\}$ , where  $R_{i,j}\in \{0,1\}$ . When given a matrix  $R$ , the  $i$ -th stage of computation is a subsequence of the original computation order, defined by  $(n_j\mid R_{i,j} = 1)$ . The recomputation plan is then determined as the concatenation of  $T$  stages. The maximum length of the recomputation sequence is  $O(T^{2})$  as a result of this formulation.
The objective of Checkmate is to minimize  $\sum_{i=1}^{T}\sum_{j=1}^{T}R_{i,j}c(n_j)$ , where  $c(n_j)$  represents the cost of recomputing node  $n_j$ . Additionally, Checkmate uses binary variables to determine whether to preserve the outputs of  $n_j$  from stage  $i$  to  $i+1$  and whether to free the outputs of  $n_j$  in stage  $i$  after evaluating  $n_k$ . The number of binary variables is  $O(|\mathcal{N}|^2 + |\mathcal{N}|M)$ , where  $M$  is the number of arcs in the DAG. Checkmate also introduces variables to represent the memory usage during the computation of node  $n_j$  at each stage. Lastly, the memory budget constraint and the dependency constraints among the variables are added to formulate the recomputation problem as a mixed-integer linear program (MILP).
A.1 Limitations of Checkmate
Despite the ability of Checkmate to consider node and value weights for any computation graph, it has several limitations.
Fixed topological order One such limitation is the requirement for a fixed topological order. When a computational graph has multiple possible computation orders (i.e., 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 to 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 is relaxed to linear programming (LP). 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 assignments for binary variables. However, Checkmate then applies randomized rounding on the binary variables that determine whether the outputs of  $n_j$  at stage  $i$  is kept for stage  $i + 1$ , which can increase the objective value. The original paper suggests setting  $0.9 \times$  budget for the LP relaxed problem, considering the memory usage increase after randomized rounding. While often there is no significant degradation in memory consumption, execution time may unexpectedly increase due to this process. We have observed that since the number of elements in  $\{R_{i,j}\}$  is  $O(|\mathcal{N}|^2)$ , it is less likely to have short recomputation sequences via randomized rounding when  $|\mathcal{N}|$  is large. The degradation in the recomputation cost is more likely to occur when the LP solver's tolerances are not small enough. However, determining the solver tolerances becomes a trade-off between convergence condition and optimization time.
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 applied.
The original description in Section 2 considers only the memory usage of input and output values. This setting is reasonable for two reasons: first, deep learning frameworks and neural network compilers have shape propagation, making it easy to infer the sizes of intermediate tensors. Second, in 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 memory usage reaches its peak, we must consider node-intermediate memory usage; otherwise, the estimation error will be significant. Additionally, certain operators, such as matrix multiplication, temporal values may consume a 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 node-intermediate memory usage when it involves a 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 removed 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'(n)$  at time  $t$ , where  $s'(n)$  denotes the intermediate memory usage of  $n$ .
B.2 Memory minimizing topological ordering
Our algorithm can also be used to find a memory-minimizing computation order without recomputation. As mentioned in Appendix A.1, 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 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. Offloading involves moving data to secondary memory and back when it is 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  $S = (n_{1},\ldots ,n_{T})$  and fetching sequence  $(F_{1},\dots ,F_{T})$ , where  $F_{i}\subseteq \mathrm{input}(n_{i})$ , that minimize the objective function. Here,  $F_{i}$  denotes which input values are fetched (instead of using previously produced ones) when computing  $n_i$ . If all  $F_{i}$  are set to empty, the problem is 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  $\sum_{i = 1}^{T}(c(n_{i}) + \sum_{v\in F_{i}}c^{\prime}(v)) + \sum_{v\in \{v|\exists i.v\in F_{i}\}}c^{\prime \prime}(v)$ , where  $c^{\prime}(v)$  and  $c''(v)$  denote the cost of fetching and offloading, respectively. This objective function assumes that offloading and fetching 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, as the optimal timings of offloading can be determined when they are given.
Extensions in algorithm We initialize all  $F_{i}$  to be empty and allow an additional mutation in each SA iteration.
4. (Flip fetch) Choose a random time  $t$  such that  $n_t \neq \mathrm{nop}$  and random value  $v \in \mathrm{input}(n_t)$ . If  $v \in F_t$ , update  $F_t \gets F_t \setminus \{v\}$ . Otherwise, update  $F_t \gets F_t \cup \{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
Figure 5: The update of lifetime by node addition or node removal. By adding a recomputation node, the life intervals of X and Z change. For either case, the update of 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 is 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 can optimize recomputation since the last value will be always recomputed from  $X$ .