SlowGuess/NeurIPS_2023_OCR huggingface dataset:
of  $v$  can be done in the same way as other mutations by considering the fetching and offloading as the addition or removal of a special node which outputs  $v$  from empty inputs.
C Technique for fast and better convergence
This section presents various techniques to expedite the convergence of SA or obtain better solutions.
C.1 Node grouping
As mentioned in Section 2.2, node grouping is a crucial pre-processing step for recomputing a sequence of nodes. Without node grouping, the SA has small chances of reaching an solution that effectively reduces the cost of the target function.
Figure 6 provides an example of node grouping. Even without the SA, node grouping can be used as a simple recomputation strategy to lower peak memory usage. In practice, our implementation uses the following algorithm. The  $\text{do\_fuse}(n_i)$  function in Algorithm 1 determines whether to fuse
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 (1) and the memory budget  $6.7\mathrm{e} + 10$ , which is  $10\%$  of the original memory consumption.
a node  $n_i$  with its succeeding operators. We perform the fusion if the outputs of  $n_i$  can be computed from inputs with a smaller memory footprint, i.e.,  $\sum_{v \in \mathrm{input}(n_i)} s(v) \leq \sum_{v \in \mathrm{output}(n_i)} s(v)$ .
Algorithm 1 Node grouping
Require: Valid computational sequence  $S = (n_{1},\ldots ,n_{T})$    
Ensure:  $S$  is a valid sequence with grouped nodes. 1: for  $i = 1,\dots ,T$  do 2: if do Fuse  $(n_i)$  then 3: for  $j = i + 1,\dots ,T$  do 4: if output  $(n_i)\subseteq$  input  $(n_j)$  then 5:  $n_j\gets$  grouped node of  $n_i$  and  $n_j$  . 6: end if 7: end for 8:  $n_i\gets$  nop 9: end if 10: end for
The node grouping algorithm with the do_fusion function mentioned above has several characteristics. Firstly, Algorithm 1 does not depend on the initial topological order of the computational graph and therefore works even if the original sequence is not memory-efficient 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 recomputation techniques by Chen et al. [7]. In fact, the above node grouping algorithm can reduce the peak memory of such graphs to a constant size providing that the all the values have the same size.
Figure 7 demonstrates the importance of node grouping in recomputation for LLaMA. When node grouping was not used, the SA converged to a 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), we were able to reduce memory usage for LLaMA to  $7.8\mathrm{e} + 10$ , which is close to the  $6.7\mathrm{e} + 10$  budget, and lower than the optimized memory usage of  $3.5\mathrm{e} + 11$  achieved without node grouping.
C.2 LogAddExp segment tree
In our optimization approach, we assume that the objective function is a function of peak memory  $M(S)$  and computational cost  $C(S)$ . For each mutation in the SA, the cost may be updated, but
Figure 8: SA convergence for ViT model. We ran the first SA on grouped nodes for 20M iterations and the second SA on decomposed nodes for 2M iterations.
the peak memory  $M(S)$  changes only when the mutation affects the peak timings. Suppose there is a computational graph with two memory peaks. In this case,  $M(S)$  is determined by the higher of the two peaks, and mutations that reduce the memory usage around the second peak do not change  $M(S)$ . In general, multiple local memory peaks make it challenging to effectively reduce the objective value.
We can overcome this issue by employing range-add/range-logaddexp segment trees instead of range-add/range-max segment trees, introduced in Section 2. The range-add/range-max segment tree recursively calculates  $\max (a,b)$  for a given interval and determines the maximum value within a range  $[l,r)$ . On the other hand, the range-add/range-logaddexp segment tree uses the recursive calculation of  $\log \mathrm{addexp}(a,b)\coloneqq \log (\exp (a) + \exp (b))$  instead of taking the maximum. The logaddexp segment tree takes into account multiple memory peaks, enabling faster convergence of the SA.
Advantages and disadvantages While the logaddexp-based segment tree is expected to have better 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 an appropriate base for log and exp. In the implementation used for Figure 8, 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  $|\mathcal{N}|$  is large. There are important considerations to be made regarding the values held by the segment tree. In a range-add/rangemax segment tree, the  $t$ -the element of the segment tree represents the memory consumption at time  $t$ . However, in 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,  $\log \text{addexp}(x,x) = \log 2 + x > x$ , meaning that repeatedly calculating  $\log \text{addexp}(a,b)$  for values  $a$  and  $b$  close to zero will result in increasingly large values. In our implementation, we utilized the fact that the peak memory is only achieved for time  $t$  where  $n_t \neq \mathrm{nop}$ . We initialized the elements of the segment tree with a large negative constant and added a large constant again when a node is added at that time, ensuring more precise estimation of the peak memory.
Figure 8 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.
Dynamic segment tree The dynamic segment tree is a data structure that can efficiently handle the following operations: (1) Insert a value at index  $i$ . (2) Remove the value at index  $i$ . (3) Add value  $x$  for range  $[l, r)$ . (4) Calculate the maximum value for range  $[l, r)$ . This kind of data structure can be created by extending binary search trees (for details, refer to textbooks such as [9]). By using this segment tree, we do not need nop, introduced in Section 2.2, and can optimize a computation sequence consisting of only  $\mathcal{N}$ . However, the dynamic segment tree often involves heavy implementation, so 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  $T$  is the maximum length of the computational sequence during optimization.
In this section, additional information for the experiments done throughout the paper is provided together 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 PyTorch's symbolic tracing mechanism, which is available for PyTorch 2.0 or later. For our experiments, we used PyTorch 2.1.0.dev20230404 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_module. This module acts as a wrapper of a 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. It then combines them into a joint graph before optimizing and partitioning it for a new model with optimized forward and backward computations. Users can pass a custom partition function (an optional argument) to aot_module to take over the optimization process.
We implemented a custom partition function that follows these steps: it receives the joint graph and applies the recomputation algorithm to create a new computation graph that produces the same 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 it into forward and backward passes as it is more memory-efficient. 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.
Preprocess for PyTorch intermediate graph representation We perform Dead Code Elimination (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 reasons: dead code can affect node grouping performance and Checkmate does not support DCE.
Next, we convert the PyTorch intermediate graph representation (torch.fx. Graph) to an internal graph format with explicit nodes  $\mathcal{N}$  and values  $\nu$ . 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. Graph, it is represented as a single node that outputs a 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 torch.fx. Graph to a more general graph format to handle operators with multiple outputs efficiently. After the recomputation is done, we convert our 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.
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] to solve large-scale LP problems in a multithreaded manner. PDLP was the fastest solver among the available options (CLP, GLOP, and SCIP) in OR-Tools.
The execution time and solution quality of Checkmate LP are highly dependent on the hyperparameters of the LP solver, as discussed in Appendix A.1. We adopted the default values of 1e-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 as the final output after testing these thresholds.
D.3 Environment
The proposed algorithm (FastSA) and Checkmate LP were evaluated using a cluster system, with each instance configured with 8 CPU cores (Intel(R) Xeon(R) Platinum 8380 @ 2.30GHz) and 100 GiB of RAM with a NVIDIA A100 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 O3 option. The compiled module was integrated with PyTorch 2.1.0.dev20230404 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^{20}$  and initialize all of them by nop. Then, for each  $n_i \in \mathcal{N}$  of the default computation sequence  $(n_1, \ldots, n_k)$ , we set  $S[t] \gets n_i$ , where  $t = \frac{i}{k + 1} T$ .
First-stage SA The SA algorithm in the first stage is applied to grouped nodes for efficient memory 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 if 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\%$  of the initial temperature. The temperature in the  $i$ -th iteration is calculated as  $h_i \times \exp(\log(h_f / h_i) \times i / N)$ , where  $h_i$  and  $h_f$  denote the initial and final temperature, respectively, and  $N = 2 \times 10^7$  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 are 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 3, we set batchsize as 512. For GPU benchmarks, we reduced the batchsize to 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 4, we used batchsize 128 and context length 512 to obtain the simulated memory and computational cost after applying recomputation. For the real GPU benchmarking, 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 benchmarking, we reduced the values to 4 and 1024 for all of them.
name	full model name	resolution	|N|	|V|
ConvNeXt [24]	convnext_tiny	224	932	1253
ConvNeXt V2 [42]	convnextv2_large	224	2840	3467
EVA-02 [13, 37]	eva02_largePatch14_224	224	6205	6967
ViT [36, 12]	vit_largePatch16_224	224	2817	3362
ViT (small)	vit_smallPatch16_224	224	1425	1706
MobileNetV3 [18, 41]	mobilenetv3_large_100	224	484	1155
EfficientNet [38, 41]	efficientnet_b0	224	996	1766
DeiT III [39]	deit3_basePatch16_224	224	1545	1850
XCiT [1]	xcit_tiny_12_p16_224	224	2818	3550
BEiT [2, 12]	beit_basePatch16_224	224	1835	2187
CoAtNet [10]	coatnet_2_rw_224	224	2204	3128
VGG11 [35]	vgg11	224	69	119
ResNet18 [16]	resnet18	224	171	437
Table 3: List of timm models (PyTorch Image Models).  $|\mathcal{N}|$  and  $|\mathcal{V}|$  denote the numbers of the nodes and values of the computational graph.
name	full model name	|N|	|V|
ALBERT [22]	albert-base-v2	2285	2419
BERT [11]	bert-base-uncased	2078	2386
DistilBERT [31]	distilbert-base-uncased	1060	1228
ELECTRA [8]	google/electra-small-discriminator	2099	2409
Table 4: List of transformers text models.
name	full model name	|N|	|V|
GPT2 [29]	gpt2	1889	2179
GPT Neo [6] 125M	EleutherAI/gpt-neo-125m	2088	2378
GPT Neo 2.7B	EleutherAI/gpt-neo-2.7B	5488	6238
BLOOM [32] 560M	bigscience/bloom-560m	3667	4167
BLOOM 3B	bigscience/bloom-3b	4567	5187
OPT [43] 350M	facebook/opt-350m	3971	4581
OPT 6.7B	facebook/opt-6.7b	5245	6059
LLaMA [40] 7B	default config*	8330	8722
Table 5: List of large language models. *For LLaMA, we constructed a model using default configuration (7B parameters) provided by transformers.LlamaConfig.
D.6 Solution stability of FastSA and Checkmate LP
Since both FastSA and Checkmate LP are randomized algorithms, the recomputation plans depend on the seed. We executed FastSA and Checkmate LP for various memory budgets and see how smoothly the solution change. Figure 9 shows the time-space tradeoff curve for different memory 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, we set memory budgets as 0.15, 0.20, ..., 0.95, 1.00 and plotted the results. In most cases, FastSA found even better solutions than Checkmate LP. Although FastSA is a randomized heuristic, its performance is stable and not heavily dependent on random seeds. On the other hand, Checkmate LP's performance is not stable, mainly due to LP-related issues as discussed in Appendix A.1. For instance, in ViT (small), one of the solutions found by Checkmate LP had almost 10 times overhead. Although running Checkmate LP with different seeds or hyperparameters may still yield good solutions, it requires tremendous computational resources.
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 MobileNetV3, FastSA could not find solutions for small memory budgets. This would be due to the limitation of the node grouping.
Figure 10: Time-space tradeoff for vision models with  $25\%$  (darker) and  $50\%$  (lighter) memory budgets. The batch size is 512. Checkmate LP was unable to find solutions within a 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 LP for DeiT III and BEiT were 2.82 and 3.52, respectively.
D.7 Full experimental results for comparison with Checkmate LP
The graph in Figure 10 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 a 100 GiB RAM, it failed to find a feasible solution within the 6-hour time limit. Furthermore, 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 as much as 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 4 and recent language models in Table 5. 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
Figure 11: Time-space tradeoff for text models, including language models, with  $25\%$  (darker) and  $50\%$  (lighter) memory budgets. Checkmate LP failed for larger models due to out-of-memory 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 in 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 estimated values. This is because it is hard to 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 in actual GPU memory. In this section, we provide a summary of the results of actual memory usage and execution time measured on an NVIDIA A100 80GB GPU and discuss the differences between the actual metrics and the simulated ones.
Figure 12 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 memory usages were less than the simulated values by  $10\%$  for no optimization cases, but could be up to  $5\%$  more for  $25\%$  budget cases. In 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 not allocate new GPU memory. Our memory simulation does not consider this information, leading to overestimation of 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 an inaccurate memory simulation. One of the reasons is that we do not consider the node-intermediate memory usage for simulation, which underestimates the memory usage. We can address this issue in our problem setting by simulating this value (please refer to Appendix B.1).
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
Figure 12: 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 batchesizes and context lengths are modified to run the original model without recomputation for comparison.
Figure 13: 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 is indicated by the blue line.
similar simulated and actual overhead is a balanced mix of operators with high costs, such as matrix multiplication, and operators with negligible costs, such as reshape, effectively averaging the cost. By accurately profiling the operators, it is possible to obtain more efficient recomputation plans.