SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Akifumi Imanishi *
Preferred Networks imanishi@preferred.jp
Zijian Xu*  
Preferred Networks  
joe@preferred.jp
Masayuki Takagi  
Preferred Networks  
mtakagi@preferred.jp
Sixue Wang
Preferred Networks
cecilwang@preferred.jp
Emilio Castillo  
Preferred Networks  
ecastill@preferred.jp
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 [19] or Moccasin [3]. 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.
The memory requirements for deep neural networks continue to grow together with model complexity. For example, training a state-of-the-art Large Language Model (LLM) such as LLaMA [40] with 65 billion parameters needs 1024 NVIDIA A100 GPU devices for 21 days and different techniques 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 the 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. (1) It 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 further tuning. Heavy inter-device communication may result in a GPU utilization of only  $5\%$  of the hardware peak[30].
Figure 1: Optimization time for all 23 models used in our experiments with different budgets. Checkmate LP could not find a feasible solution within 6 hours for models with more than 2300 nodes. For most instances, FastSA could found recomputation plan around 3 seconds.
Numerous techniques exist to lessen memory requirements while training deep neural networks. Gradient checkpointing, or recomputation, is a widely-known approach that recomutes some of the activations during the backward pass, rather than saving them in memory, consequently reducing 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 computation 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 one 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 a 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 a 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.1). A recent work, Moccasin [3], uses constraint programming to build a set of restrictions and introduces a hyperparameter that limits the amount of times a variable can be recomputed. 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 a fast heuristic algorithm that operates on 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 finding a candidate sequence of operators that minimizes an objective function. We apply a simulated annealing algorithm where each step can be an operator addition or 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 is significantly improved by using a segment tree data structure [9] which allows to lazily evaluate the memory usage of each candidate sequence, enabling nearly 1 million mutations on 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 a 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\%$  computational cost on average.
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.
2 Fast Simulated Annealing for Recomputation
2.1 Problem Description
To apply the Simulated Annealing optimization strategy to 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 a finite set of nodes  $\mathcal{N}$  and a finite set of values  $\nu$ . Each node takes some values as inputs and produces some values as outputs; therefore, we introduce notations  $\mathrm{input}(n)\subseteq \mathcal{V}$  and  $\mathrm{output}(n)\subseteq \mathcal{V}$  for the values associated to a node  $n\in \mathcal{N}$ . In most cases, the numbers of input  $(n)$  and output  $(n)$  are small. The computational graph is often visualized as a directed acyclic graph (DAG) with directed edges from  $n_i\in \mathcal{N}$  to  $n_j\in \mathcal{N}$  if an output of  $n_i$  is used as an input of  $n_j$ , i.e.,  $\mathrm{output}(n_i)\cap \mathrm{input}(n_j)\neq \emptyset$ . A value that is not the output of any node is a 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)\geq 0$ , and the size of value  $v$  is represented by  $s(v)\geq 0$ .
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  $(n_{1},\ldots ,n_{T})$ , which means that  $n_t$  is operated at integer time  $t$ . We say that a sequence of nodes has a valid dependency if we can execute the operators in the order of the sequence so that all model outputs are correctly computed. Formally,  $\forall i.\forall v\in \mathrm{input}(n_i).\exists j < i$  such that  $v\in \mathrm{output}(n_j)$  or  $v$  is a model input. It is worth noting that node duplication is allowed in the above node sequence  $(n_1,\dots,n_T)$ , and the recomputation of a value can be represented by multiple appearances of the same node.
Memory usage In order to optimize memory usage, we must keep track of when values are computed and used as inputs. Specifically, we must keep value  $v$  in memory from the time it is computed as an output until the last time it is used as an input. However, when  $v$  is recomputed at time  $t$ , we can temporarily free  $v$  from memory if it is not used until  $t$ . When recomputation is considered, the lifetime of  $v$  can be represented as a set of independent time intervals, as shown in Figure 2. To determine whether  $v$  is kept in memory at a given time  $t$ , we use a function  $L_{S}(v,t) \mapsto \{0,1\}$ . This function outputs 1 if  $v$  is kept in memory at time  $t$ , and 0 otherwise.
$$
L _ {S} (v, t) := \left\{ \begin{array}{l l} 1 & v \text {i s a m o d e l i n p u t} \\ 1 & v \text {i s a m o d e l o u t p u t} \wedge \forall t ^ {\prime} > t, v \not \in \mathrm {o u t p u t} (n _ {t ^ {\prime}}) \\ 1 & \exists t ^ {\prime \prime} \geq t, v \in \mathrm {i n p u t} (n _ {t ^ {\prime \prime}}) \wedge \forall t ^ {\prime} \in [ t, t ^ {\prime \prime}), v \not \in \mathrm {o u t p u t} (n _ {t ^ {\prime}}) \\ 0 & \text {o t h e r w i s e} \end{array} \right.
$$
Let  $M(S)$  be the minimum memory size required to compute a valid sequence  $S$ . Then, we have  $M(S) = \max_{1 \leq t \leq T} \sum_{v \in V} L_S(v, t) \cdot s(v)$ . Also, we define the cost of the sequence,  $C(S)$ , by  $C(S) = \sum_{1 \leq t \leq T} c(n_t)$ .
Figure 3: Overview of the simulated annealing.
Recomputation Let  $f$  be an objective function which maps sequences to a real number. The recomputation problem is defined as the construction of a valid sequence  $(n_1, \ldots, n_T)$  that minimizes  $f(S)$ .
When minimizing memory usage, we set  $f(S) \coloneqq M(S)$ . The objective function can depend not only on  $M(S)$  but also on computational costs. In our experiments, we utilized Equation 1:
$$
f (S) := \max  (\text {b u d g e t}, M (S)) \times C (S) \tag {1}
$$
This function minimizes the total computational cost when the memory budget is met. Otherwise, it tries to minimize both the peak memory and the cost.
2.1.1 Add-max segment tree
The range-add/range-max segment tree [5] is a data structure that can evaluate add and sum queries for intervals lazily. The add-max segment tree, utilized in this paper, holds an array  $A[T]$  of length  $T$ , and can carry out the following operations in  $O(\log T)$  time:
- Add  $x$  to  $A[i]$  for  $i \in [l, r)$
- Obtain the maximum value of  $A[i]$  for  $i \in [l, r)$
2.2 Fast Simulated Annealing Algorithm
The proposed algorithm is based on simulated annealing (SA) with a 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 an overview of the algorithm. Two critical mutations must be performed in sublinear time for the algorithm to be efficient: (1) Slightly modifying  $S$  to create  $S'$ . (2) Computing  $f(S')$ .
To generate  $S'$  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 to a 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 no inputs or outputs associated with it. When adding or removing a node, its value is exchanged with that of nop.
We start with a valid sequence  $S = (n_{1}, \dots, n_{T})$  and for all  $1 \leq i \leq T$ ,  $n_i \in \{\mathcal{N} \cup \mathrm{nop}\}$ . To enable the 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 a random node  $n \in \mathcal{N}$  and a random time  $t$  where  $n_t = \mathrm{nop}$ . We attempt to update  $n$ , i.e.,  $n_t \gets n$ .
2. (Remove computation) Select a random time  $t$  such that  $n_t \neq \mathrm{nop}$  and attempt to update  $n_t \gets \mathrm{nop}$ .
3. (Rotating sequence) Select two random times  $t_1$  and  $t_2$  such that  $n_{t_1} \neq \mathrm{nop}$  and  $n_{t_2} = \mathrm{nop}$ . Try swapping  $n_{t_1}$  and  $n_{t_2}$ . This modification can be 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 if they are satisfied. Note that we can efficiently check these conditions by maintaining a set of produced and used times for each value  $v$ .
1. (Add computation) For each input  $v$  of  $n_t$ ,  $v$  must be produced before time  $t$ .
2. (Remove computation) For each output  $v$  of  $n_t$ , either  $v$  must already be produced before time  $t$ , or there must not be any user of  $v$  before the next production of  $n_t$ .
An important point to note is that the mutations 1 and 2 are inverses of each other, meaning that to undo a change, one simply needs to reapply the inverse mutation.
2.2.1 Updating peak memory and objective
In the following discussion, we consider the update of  $f(S)$  by mutation 1 and 2. Since the objective function  $f(S)$  often depends on both memory usage  $M(S)$  and the overall cost  $C(S) = \sum_{t} c(n_{t})$ , we assume this is still the case here. We can easily update  $C(S)$  by adding or subtracting the 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 a node is added or removed from  $S$ , we have to calculate the aggregated memory consumption by the values with overlapping lifetimes.
The lifetime of inputs and outputs change when adding or removing node  $n$  (see Figure 5 in Appendix for illustration). However, the insertion and deletion of a node only slightly modify the life intervals for each node's input and output values. Consider a value  $v$  with life intervals  $L(v) \coloneqq \{(t_i, t_j)\}$  determined by  $L_S(v, t)$ . Suppose node  $n$  is inserted at time  $t$ , and  $v$  is an input of  $n$ . We update the life intervals of  $v$  as follows: If there is  $(t_i, t_j) \in L(v)$  such that  $t_i \leq t \leq t_j$ , do nothing. Otherwise, take a  $(t_i, t_j) \in L(v)$  such that  $t_j \leq t$  is maximum (such  $t_j$  exists if the sequence after insertion is valid). Update  $L(v) \gets L(v) \backslash (t_i, t_j) \cup (t_i, t)$ . Similar rules apply for updating the life intervals of the outputs of the inserted node and the inputs and outputs of the node to remove. Specifically, on inserting or deleting a 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  $(t_L, t_R)$  to  $(t_L, t_{R'})$  or  $(t_{L'}, t_R)$ .
3. Split a range  $(t_L, t_R)$  into  $(t_L, t_{R'})$  and  $(t_{L'}, t_R)$ .
4. Merge two ranges  $(t_L, t_{R'})$  and  $(t_{L'}, t_R)$  into  $(t_L, t_R)$ .
Because the update of live intervals involves adding or subtracting the value size  $s(v)$  from a certain range for each input or output value  $v$  (as depicted in Figure 5 in Appendix), we can efficiently maintain memory usage and calculate the peak memory using a segment tree. This data structure allows for range-max and range-sum queries to be performed in  $O(\log T)$  time, as introduced in Section 2.1.1. The segment tree maintains memory usage for time  $t$  in an internal array  $A[t]$ , and peak memory can be determined efficiently by taking the maximum value in the range  $[0, T)$  in  $O(\log 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 as it can reduce the time for the differential update of a lifetime interval from the naive  $O(T)$  time to  $O(\log T)$  time.
2.2.2 Improving memory reduction by grouping
To determine an optimal recomputation sequence, it may be necessary to 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 a series of nodes. A grouping node, represented by  $g$ , is formed by concatenating two nodes,  $n_1$  and  $n_2$ , and has the following properties:
$c(g) = c(n_1) + c(n_2)$
- input  $(g) = \mathrm{input}(n_1)\cup (\mathrm{input}(n_2)\backslash \mathrm{output}(n_1))$
- output  $(g) = (\mathrm{output}(n_1)\backslash \mathrm{input}(n_2))\cup \mathrm{output}(n_2)$
Grouped nodes can be especially useful for sequenced patterns where  $n_1$  must be computed immediately before  $n_2$ , as grouping these nodes together can make the sequence more likely to 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 another round of annealing with lower temperatures. In Appendix C.1, we provide further discussion on the benefits of using node grouping for optimization.
2.3 Other Considerations
Conceptually, recomputation enables to optimize the time-space tradeoff without changing the graph semantics, i.e., 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 to reduce memory usage that can be easily incorporated into our 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 operation. 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.
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 computational 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 vision models were set up for image classification, and model variants with sequence classification 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 hyperparameters of SA As outlined in Section 2, we utilized an objective function of the form  $f(S) = \max (\mathrm{budget}, M(S)) \times 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 to 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 effectively for a broad range of practical models.
To ensure efficient convergence, we utilized the first SA on grouped nodes for at most 20 million iterations, or until the memory budget was met, whichever occurred first. The second SA ran for a fixed 2 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
In certain scenarios, this objective function may exceed the budget constraint. For experiments of Table 1 and Table 2, where memory budgets must be satisfied, we employed  $f(S) = L(\mathrm{Elu}(L^{-1}(M(S)))) \times C(S)^{0.6}$ , where  $L$  is a linear function with  $L(0) =$  budget and  $L(-100) = 0$ , determining the sharpness of Elu.
Figure 4: Comparison of simulated memory decrease and cost increase by recomputation. For each model, the memory budget is set to  $50\%$  and  $25\%$  of the simulated initial peak memory (darker 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.
could not find any feasible solution within 6 hours for all models with more than 500 nodes due to 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 (LP). For clarification, the original version is referred as Checkmate MILP. Gurobi [15] is used as the internal solver for Checkmate MILP together with PDLP solver, provided by OR-Tools [28] for Checkmate  $\mathrm{LP^3}$ .
Moccassin Another alternative to perform recomputation is Moccasin [3]. Similarly to Checkmate, 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 a value can be recomputed. Thanks to this limitation, the number of integer variables involved in the CP set is linear to the number of nodes in the graph, allowing a 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 a 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.