Coop: Memory is not a Commodity
Authors: Jianhao Zhang
Shihan Ma
Peihong Liu
Jinhui Yuan
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Tensor Rematerialization; Gradient Checkpointing; Activation Recomputing; Deep Learning; Deep Learning Frameworks; Memory Allocator

Abstract: Tensor rematerialization allows the training of deep neural networks (DNNs) under limited memory budgets by checkpointing the models and recomputing the evicted tensors as needed. However, the existing tensor rematerialization techniques overlook the memory system in deep learning frameworks and implicitly assume that free memory blocks at different addresses are identical. Under this flawed assumption, discontiguous tensors are evicted, among which some are not used to allocate the new tensor. This leads to severe memory fragmentation and increases the cost of potential rematerializations.

To address this issue, we propose to evict tensors within a sliding window to ensure all evictions are contiguous and are immediately used. Furthermore, we proposed cheap tensor partitioning and recomputable in-place to further reduce the rematerialization cost by optimizing the tensor allocation.

We named our method Coop as it is a co-optimization of tensor allocation and tensor rematerialization. We evaluated Coop on eight representative DNNs. The experimental results demonstrate that Coop achieves up to $2\times$ memory saving and hugely reduces compute overhead, search latency, and memory fragmentation compared to the state-of-the-art baselines.

Coop: Memory is not a Commodity
Jianhao Zhang"
Shihan Ma'
OneFlow Research daquexian566ggmai.cco
OneFlow Research mmasss ciiasss1050@gail1 com
Peihong Liu OneFlow Research peihong www.t1oo..ooooooooooooooooo.. com
Jinhui Yuan OneFlow Research yuanjinhui wtwphnhuuwooonffoo.oooooo org
Abstract
Tensor rematerialization allows the training of deep neural networks (DNNs) under limited memory budgets by checkpointing the models and recomputing the evicted tensors as needed. However, the existing tensor rematerialization techniques overlook the memory system in deep learning frameworks and implicitly assume that free memory blocks at different addresses are identical. Under this flawed assumption, discontiguous tensors are evicted, among which some are not used to allocate the new tensor. This leads severe memory fragmentation and increases the cost of potential rematerializations To address this issue, we propose to evict tensors within sliding window ensure all evictions contiguous and are immediately used. Furthermore, we proposed cheap tensor partitioning and recomputable in-place further reduce the rematerialization cost by optimizing the tensor allocation. We named our method Coop as it is a co-optimization of tensor allocation and tensor rematerialization. We evaluated Coop on eight representative DNNs. The experimental results demonstrate that Coop achieves up 2 memory saving and hugely reduces compute overhead, search latency, and memory fragmentation compared to the state-of-the ttthhearr art baselines.
Introduction
Recent development of deep neural networks (DNNs) shows continuous rage on increasing the scale of the network structures, which dramatically improves the capability of the neural network 2 Training such gigantic models with billions of parameters, however, requires huge on-device memory 5 Tensor rematerialization, also known as activation checkpointing. one technique that allows the training of large DNNs under limited GPU memory budget without compromising the model's accuracy. This technique has been widely applied in the training large language models such as GPT-3
6 and LLaMa 33.
The key idea of tensor rematerialization
evict some intermediate tensors during forward propaga- tion and recompute them as needed during backward propagation, essentially trading recomputing time for memory. Tensor rematerialization can be traced back to checkpointing in reverse-mon mode automatic differentiation 7 8, 99, and was first applied to DNN training as a global optimization
technique for static computation graphs (SCGs) I0O 12 13, Chen et al. (II) reduced the memory cost of training neural network with n. layers to O(VR) by dividing the network into segments and dropping the intermediate results within each segment. Checkmate (IT) further for- malized tensor rematerialization find an optimal solution. Tensor rematerialization methods that support dynamic computational
mixed integer linear program and used numerical solvers to
Equal contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

graphs (DCGs) emerge with the advancement of frameworks that utilize DCGs (e.g. Pytorch [14.. Megatron-LM [5] uses selective activation recomputation to drop activations that take up large memories but are cheap to recompute in Transformer layers. Dynamic tensor rematerialization (DTR) [6 automatically optimizes memory usage during the training of DCGs by repeatedly searching for and evicting optimal tensor until the next tensor can be successfully allocated.
However, all of the above tensor rematerialization methods have a similar problem, that is, they assume the memory in deep learning (DL) system is a fungible commodity where free memory blocks different addresses are identical. Under this assumption, evicting tensors with size x sufficient for allocating a tensor with size I, which is incorrect unless all of the evicted tensors are contiguous. Evicting tensors that do not contribute to a contiguous memory will lead to memory fragmentation. subsequently increasing the cost associated with any potential rematerialization. This motivates us to ask the question: allocating a new tensor while the recomputing cost is the minimum?
possible to release a contiguous memory block that is large enough for
To answer this question, we proposed Coop, gradient checkpointing scheme that co-optimizes tensor rematerialization and tensor allocation. To the best of our knowledge, Coop is the only tensor rematerialization scheme that fully bypasses the incorrect assumption of DL memory system. We propose that straightforward solution to evicting a set of contiguous tensors is to use a sliding window algorithm. The window is moved within the memory pool until an optimal solution is found. Evicting the tensors in the optimal window will free enough memory while the recomputing cost is the minimum. In addition, unlike prior works that do not optimize tensor allocations, we introduce smart memory allocator to further reduce the cost of evicting tensors in the sliding window. The smart memory allocator (1) groups "cheap" tensors together and (2) enables recomputable in in-place to prevent the disruption of contiguous memory blocks. We show that by co-optimizing tensor rematerialization and memory allocation, Coop reduces memory fragmentation and supports training large-scale DNNs under the most limited memory budget with less compute overhead and lower search latency. For example, when training 2.7B-parameter GPT-3 style model, Coop reduces the memory fragmentation rate and compute overhead by up to ~ 5x, the search latency by ~ 10x, and the minimum memory budget by 25% compared with DTR. The proposed strategy is -agnostic and can be universally used to checkpoint any dynamic models in real time. In summary, Coop makes the following contributions:
We argued for the first time that existing tensor rematerialization methods overlook the memory system during optimization and wrongly assume that the memory in DL systems is fungible commodity. We addressed this problem by co-optimizing tensor rematerialization and tensor allocation, which opens up a new perspective on improving checkpointing techniques in DL frameworks.
Coop allows the tensor rematerialization scheme and memory allocator to facilitate each other. The heuristic of tensor rematerialization reduced by taking into account the properties of memory allocators, and the memory allocators are improved by considering the efficiency of different operations in tensor rematerialization.
We demonstrated that Coop enables the training of the most widely-used DNNs with less compute overhead, lower minimum memory budget, and smaller search latency.
Coop has been implemented in OneFlow framework and it can be easily integrated into any other deep learning framework.
2 Background and Motivation
2.1 Memory Allocator in DL Systems
Memory allocator is an important component of deep learning systems. All known DL systems (e.g. PyTorch (e.g.PyTorch 14, TensorFlow [17, and MXNet (18) have their own memory allocators to enable fine-grained memory management and avoid the overhead of communication with the operating system (OS). Compared advanced CPU memory allocators, e.g. mimalloc
[9 and jemalloc [20]. memory allocators in DL frameworks are simpler. Typically, in DL frameworks, when a tensor is destructed, its memory not returned to the OS but inserted into the free list of the allocator and merged with other chunks with adjacent memory addresses. When the user requests memory for tensor with size S, the allocator tries to place the tensor at the leftmost side of a free chunk with a size equal to or larger than S. If the sizes of the free chunks are all less than S, the allocator will

---

Coop
optimal tensor
Sliding
(h)
Evict an optinal lossor
Recomputable in-paaee
(0)
Repeat until enough free memory
Chesp tensor portitioning
"cheap"
"expessiive"
unevietable
empty
ovictions
Figure 1: Comparison between DTR and Coop. DTR overlooks the underlying memory system, which leads to redundant evictions. In contrast, Coop co-optimizes tensor rematerialization and tensor allocation find the best tensors to evict. (a) A typical memory layout given by the conventional tensor allocator, where all tensors are mixed together. Tensors are classified by their cost densities (computational cost divided by memory size). The "cheap" and "expensive" tensors denote the tensors with low and high cost densities, respectively. Unevictable tensors, including network parameters and buffers, cannot be evicted. (b) Coop uses the sliding window algorithm find the optimal set of tensors under given memory layout. (c) Coop uses recomputable in 1e-place operations to optimize memory layout by ensuring that the unevictable tensors remain in the same place even after being updated. This optimization reduces the cost of the tensors in the sliding window compared to (b). (d) With cheap tensor partitioning. "cheap" and "expensive" tensors are separately allocated to the two sides of the memory pool. The memory layout is further optimized with the eviction cost reduced.
request new memory from the OS, even though the total size of the free chunks might be greater than S. These unused free chunks are called memory fragments. During training, all intermediate activations and gradients are released at the end of each iteration, which produces large free memory chunks. Therefore, tensors generated by sequential operations tend be allocated to contiguous spaces from the leftmost of the memory pool in each iteration.
Besides, it is possible for memory allocators leverage the page table of the hardware and re-compose discontinuous memory chunks into a chunk continuous from the virtual memory's 's view. Studies have proven that memory fragmentation in CPU be reduced in this way (2 22] However, the driver of NVIDIA GPUs is a proprietary NVIDIA-controlled black box, making it almost impossible for community researchers apply similar ideas to GPUs.
2.2 Tensor Rematerialization and Memory Fragmentation
Tensor rematerialization in DNN training is first introduced by Chen et al. [10]. They divide the network into segments, evict all intermediate features within each segment, and recompute them during backpropagation. Subsequent studies enhance this method by intelligently selecting tensors to evict. Dynamic Tensor Rematerialization (DTR) [6 devises heuristic guide the eviction of tensors during runtime. The heuristic leverages the metadata of each tensor (staleness, memory, and computational cost) to select the stalest, largest, and cheapest tensor to evict when the memory runs out. This process runs several times until the released memory is sufficient for the new tensor. Recent work MegTaiChi TaiChi [23] proposed Dynamic Tensor Evicting (DTE) based on DTR. DTE modifics DTR's heuristic by considering the adjacent free memory blocks of each tensor.
These tensor rematerialization methods enable the training of large DNNs with limited memory budgets and low overhead. However, their performance is inhibited by large memory fragmentations. As a real-world example, DELTA (24] implements DTR Parrots framework (25) and reports that the memory fragmentation hugely affects the maximum capacity of the model. Memory fragmentation the small and sparsely scattered memory chunks that are unusable for tensor allocation. This problem affects all DNN training but is exaggerated by tensor rematerialization, which incurs a more frequent rearrangement of the memory. A fundamental reason why the existing tensor rematerialization methods cannot address this problem is that they do not emphasize producing
3

---

contiguous chunk of free memory. Additionally, their heuristics render the condition worse. Most rematerialization methods penalize evicting tensors generated by sequential operations reduce recursive rematerializations. For example, the projected cost of a tensor in DTR's heuristic is calculated over the tensor's 's neighborhood (tensors that depend on or are relied on recomputing the current tensor). Penalizing evictions of sequential tensors often results in discontiguous free memory chunks, as the sequential tensors tend to locate contiguously (see 2.1]. To address the memory fragmentation in tensor rematerialization, Coop proposes to evict a set of contiguous tensors by using a sliding window algorithm, which is hugely different from the modification of heuristic as in DTE.
2.3 In-place Mutation in DL Systems
In-place mutation is a common and important feature in deep learning frameworks. For ex- ample, PyTorch users often construct their networks with in-place mutating operations like nn. Me.RRLU (ipaaceele)))))) tensor zero_ O, torch. add (a, b, out=c), avoid allocating memory for the output tensor by directly overwriting the existing tensor. In-place operations are more friendfy and use less memory than their out-of-place counterparts. The in-place operations produce correct results as long as the original values are not used in the future. In deep learning frameworks, this requirement met by not capturing the original value during backward propagation. However, in frameworks with tensor rematerialization, the original value is likely to be used during recomputation, which leads to incorrect results. For example, if the user runs two operations y x + 1 and relu_ () in sequence and y needs to be rematerialized afterward, the operation x + will run again but the original value of has already been discarded by x. relu O)
DTR solves this problem by introducing a copy-on-write layer, which copies the tensor and mutates the copied one. This actually transforms the in-place operations into out-of-place operations, which negates the ad advantages of in-place mutations and may cause potential problems (explained in Section B.5). Recently, Koka (26). a functional programming language, proposed novel programming paradigm called functional but in- -lape (FBIP) in their new memory management system Perceus (27]. FBIP enables in- in-place mutation in a purely functional manner by allocating new objects the memory of known unused objects. Inspired by FBIP, we proposed recomputable in-place module in Coop to prevent unnecessary evictions and reduce memory fragmentation.
2.4 Other Memory-saving Techniques
In addition to tensor rematerialization, several other techniques can be used to reduce the mem- ory requirements for training large models, such as model parallelism [28. gradient accumulation, reversible operations 229 30), quantization [31], and offloading 6 32). These methods are com- plementary to tensor rematerialization. The combined use of these techniques can result in more significant memory savings 3 as demonstrated Turing-NLG, GPT-3, and LLaMa 6 3.
Nonetheless, each of these techniques comes with its own limitations. Model parallelism generates additional communication overhead (34]. Gradient accumulation necessitates reduction in batch size, which can adversely affect training accuracy, particularly for models with batch normalizations 35 36 Reversible operations do not possess universal applicability. Quantization may entail lossy compression [37]. Offloading is typically slower than rematerialization and requires prefetching for satisfactory performance (33] making challenging to implement in DCGs.
3 Methods
3.1
Problem Formulation
Coop is a runtime memory manager that allows the training of dynamic DNNs under a preset memory budget by manipulating the state of the memory pool. The state of the memory pool is denoted as memory layout, including the position and the size of each tensor. Coop manipulates the memory layout by intercepting tensor allocation, eviction, and rematerialization.
Assume that there are N tensors in the memory pool after several operations in the forward process. Each tensor has memory size mll) and heuristic h(l). Heuristic is the cost of evicting tensor (defined in 3.3), the less the better. Coop aims to find an optimal set of tensors to evict when the memory meets the budget, such that (L) evicting the set of tensors is sufficient for allocating the new

---

lensor and (2) the evictions bring in the minimum potential rematerialization cost. Guided by this target, we introduce the cost function Coop as follows:
arg min S,L
Σ h(t),, subject to M(S. 2 MR
(I)
FES
where denotes a memory layout, S denotes a set of tensors, M(S, L) is the largest available contiguous memory under memory layout L after evicting all tensors in S, and Mg is the required memory size.
3.2 Method Overview
We observed two desiderata from the cost function, which inspired to propose the three modules in Coop. First, free but not contiguous memory blocks in S do not contribute to the largest contiguous memory M (S, T.) and thus will not be used for allocating new tensors. In other words, it is more efficient to evict tensors that produce a single contiguous block. We used a sliding window algorithm to address this problem in ((N) time complexity (.3) Second, memory layout I. affects the cost of evicting tensors in S. and L is related to lensor allocation in DCGs. Given that contiguous block can be produced by the sliding window algorithm, we further introduced two modules to reduce the cost of evicting tensors in this block by managing the memory layout: (1) we proposed cheap tensor partitioning 3.4 to cluster tensors that are computationally "cheap" to the same place. (2) we introduced recomputable in-place in 3.5 to avoid the random reordering of the memory layout that splits contiguous memory, An overview of Coop is displayed in Figure and the algorithm of Coop is shown in Algorithm]
Algorithm 1 The algorithm of Coop
procedure ALLOCATE(OP, size)
op is an in-place operation then addr +- input.addr
else
block +- free_free____ sece__bock__eee less_ess___ than(size)
block is None then
evict( (riddisg window search( (szze))
block < find_free_lee block_ook_______ allok_noo__n l
if op is expensive then
addr +- le ft__addr
else
addr + block. right _addr- - size
end if
end if
end if
return addr end procedure
3.3 Sliding Window Algorithm
brute-force approach to solving Equation is to traverse the combinations of tensors and find a valid combination with the minimum cost. It requires total of 2 enumerations (N is the number of tensors in the memory pool), which is computationally unfcasible. To circumvent this issue, we used the sliding window algorithm to find the optimal solution and reduced the time complexity from O(2N) to ONN..
We track the states of all tensors by storing them in list, sorted by the memory addresses. The free memory chunks are included as "special tensors" of which the heuristics are zero (no cost to recompute them). In this way, finding a set of tensors that produce a single contiguous block is equal to finding a subsequence of the list. We use two pointers to denote the start and the end of the sliding window. Tensors in the window are supposed to be evicted. By moving the window within the list and continuously comparing the summed heuristics of the tensors in the window, a set of contiguous tensors of which the summed memory is larger than required and the cost minimum can be found.
5

---

To calculate the eviction cost, we need to choose a heuristic that is consistent with our sliding window e(t)
algorithm. We define the heuristic h(t) s(t) where c(t) is the projected cost (the sum of the computational cost of tensor t and t's evicted neighborhood) and s(t) is the staleness. Different from c(t)
DTR's 's heuristic (h(t) (()))=a(1)) where m(t) is the memory size of t), our heuristic does not need the information of tensor's s memory size. The reason is that the information of tensor memory has been embedded into the layout S, thus the list maintained by Coop. The sliding window algorithm utilizes this information and searches for an exact set of contiguous tensors to evict rather than repeatedly and blindly looking for tensors that might be large enough to cover the new tensor.
One benefit of the proposed sliding window algorithm is that traversing for one time gives the best tensors to evict. Therefore, the time complexity for searching is O(N). Another particular advantage is that no tensors are "innocently" evicted, as the freed memory after each eviction is contiguous and meets the required memory exactly in the value. By contrast, DTR evicts one tensor each time, resulting in sparsely distributed free memory blocks that cause the memory fragmentation issue.
3.4 Cheap Tensor Partitioning
Given that contiguous free memory chunk can be produced by the proposed sliding window algorithm and the overall heuristic is the minimum under the current memory layout, we ask whether is possible to optimize the memory layout as well to further reduce the overall heuristic. We propose an approach called cheap lensor partitioning to address this problem. Instead of allocating tensors contiguously as in the traditional DL systems, cost to the same place. Clustering tensors with the same magnitude reduces the overall heuristic in two ways: (1) it breaks the pattern that tensors generated by sequential operations are also allocated contiguously in memory, thereby reducing the projected cost c(t) in the heuristic; (2) it increases the chance to evict a set of contiguous "cheap" tensors. In this way, proper memory layout is prepared before tensor eviction.
group tensors with the same magnitude of eviction
DTR
X
X1
X2
X3
X1
1x1, 64
X3
Xg
eviet
relu
X2
xs
allocute Xs
3x3, 64
Coop
Xy
relu
Xp
X₂
Xg
X4
X4
Xg
evict xp, X2
1x1,128
X5
X3 allocute Xg
Figure Illustration cheap tensor partitioning in Coop. (a) Layers of a representative convolutional neural network, with convolutions followed by activation layers. (b) Tensor eviction under DTR with tensors allocated in sequence. (c) Tensor eviction under Coop with cheap tensor partitioning. Tensors are allocated from both sides of the memory pool.
We implemented cheap tensor partitioning by allocating tensors from the leftmost and rightmost of the memory pool (Figure (c)), which is feasible because the memory budget is given by the user before training. Cost density (computational cost divided by memory size) is used to evaluate the magnitude of eviction cost. Tensors with the same magnitude of cost density are allocated to the same end of memory pool during the forward propagation of the model. We observed that most of the operators in NN can be simply classified into two categories by their complexity: super-linear (e.g. matmul and conv, denoted as C1) and linear/sub-linee (e.g. element-wise ops, denoted as C2). For most cases, the cost density of super-lineat operations is larger than the other operators by one order magnitude. Cost densities of major operators in DNNs on NVIDIA GeForce RTX 2080 Ti summarized in Table
Figure a shows an example of a typical convolutional neural network, which each convolution
layer is followed by activation function. Here we did not include Batch Normalization as it has

---

Table 1: Cost density of major operations in DNNs (ps/MB).
Operation
ResNet-500
GPT-2
U-Net"
Swin-T'
C₁ Conw/MatMul 35.6
33.5
89.3
32.7
C2 Batch/LayerNorm
5.0
4.2
5.3
4.1
C. ReLU/GELU
3,9
3.8
3,9
3,9
Model with Conv, BatchNorm and ReLU Model with MatMul, LayerNorm and GELU.
similar computational cost density with the activation layers [T and is not necessarily included all neural networks. Suppose the tensors Zo: .... I4 are generated at the start of training, thus stored sequentially in memory. If the memory is full and another 100 MB for Fı is required, the two tensors generated by the two activation layers (Ny and Z2, each 50 MB) will be first evicted under DTR's strategy (ru is the stalest and cheapest tensor in the memory and its eviction increases the heuristic of 21). However, as the freed chunks are not contiguous. extra tensors are required to be freed, e.g. 21. bringing in useless evictions and leading to memory fragmentation (Figure (b)..
3.5 Recomputable In-place
When dealing with operations that directly mutate the content of the input tensor, i.e., in in-place operations, DTR introduces copy-on- write layer to replace the in-place operation with three separate operations: copying the input tensor, mutating the copied tensor, and evicting the input tensor (Figure (c)). Allocating
memory for the copied tensor has detrimental effects: (1) if the memory pool is full, additional tensors are required to be evicted for allocating the newly copied tensor, which takes time and brings extra recomputation cost; (2) more critically, this tensor unevictable, such as the parameter that is updated in- in-place after each iteration, a random allocation of this tensor will partition the memory into non-contiguous parts, preventing the creation of contiguous free memory.
(a) Initial State (b) Naive
(c) DTR (copy-on-write)
Coop (recomputable in-place)
J
G
evicted
shared mem o 0f
evieted
la al
relu
P xrellul
new nemm
new mom
reused
Figure 3: Illustration of different in-place mechanisms on relu (in-place ReLU), (a) Initial state of the operations. (b) Unrecomputable naive in-place operations. The dependencies between adjacent variables are incorrect. (c) Copy-on cepyy-o-wiieg write mutation in DTR. (d) Recomputable in-place in Coop that avoids allocation of new tensors.
Inspired by the functional but in-place paradigm in Perceus (27, Coop proposes a memory-reuss mechanism named recomputable in-place to handle in-place operations and further optimize the memory layout. Instead of copying the input tensor to a new one, we reuse the memory of the input tensor and directly allocate the new tensor to its original place such that the input tensor and the new output tensor share the same memory (Figure 3 (d)). As the input tensor is evicted after the in-place operation, recomputing tensors that rely it requires recomputing the input tensor itself first. Therefore, the value of the input tensor is always the same as its initial version.
Besides, we observed that the in-place operations mostly happen during the updates of model parameters that are unevictable. To prevent these tensors from partitioning the whole memory pool into discontiguous chunks, Coop first allocates the tensors of parameters to the two ends of the memory pool, and then reuses their memory during the updates (Figure T()). In this way, the rate of memory fragmentation can be reduced to the least.
7

---

Evaluation
4.1 Evaluation Setup
We compared our proposed method, Coop. with two state-of-the-art DCG-based automatic rema- terialization methods, DTR [66] and DTE (the rematerialization method introduced in MegTaichi
23 ). across eight representative DNNs. The eight DNNs include 2.7B parameter GPT-3 style large language model, Sriin-Trannformer ResNet-50, etc. Additionally, we compared our method with selective activation recomputation [38] (the rematerialization method designed for Transformer models in Megatron-LMMe LM on Transformer models. Comparisons with static graph methods are
beyond the scope of this paper due to the distinct prerequisites and use cases.
Table 2: Rematerialization methods compared in evaluation.
Method
Description
Automatic MS -aware
MS- ASSoopimmzed
Traversal Count
Coop
The proposed tensor rematerial-
Single
ization method.
DTE
Our impl. of Dynamic Tensor
X
Multiple
Evicting (23 in OneFlow.
DTR
Our impl. of Dynamic Tensor Re-
X
Multiple
materialization [16] in OneFlow.
SAR
Our impl. of Selective Activation
X
X
X
Recomputation 38
OneFlow.
MS is the abbreviation the memory system.
DTE's heuristic considers adjacent free memory but cannot promise contiguous memory block obtained.
Table 2 provides an overview of the methods compared in our evaluation. The baseline methods (DTR, DTE, and SAR) have been implemented in different frameworks. For a fair comparison, we re-implemented all baselines in OneFlow (39), which is an open-source deep learning framework with PyTorch-lligned APIs. In our DTE implementation, we only consider the cost of recomputation, whereas the original paper also considers the cost of swapping tensors between GPUs and the host 23 We set the coefficient of computing times in DTE to align with DTR and Coop. We also compared Coop in OneFlow with the official DTR in Py Torch and the official DTE in MegEngine in Appendix A. The effects of the three proposed modules were examined Appendix B. All experiments were conducted on a machine equipped with 4 NVIDIA A A100 GPU (80 GB, CUDA 11.7, CuDNN 8.5.0) and 56 Intel(R) Xeon(R) Platinum 8336C CPU cores running Ubuntu 20.04. For BERT Large and GPT-3 style 2.7B, Adam optimizer was used, while the SGD optimizer was used for the other experiments. ZeRO stage 2 [40 is used when training the GPT-3 style 2.7B model. Among the eight DNNs, BiLSTM and SPOS have dynamic network structures that vary based on the input.
4.2 Evaluation Metric
Five metries are used for evaluating the performance of the methods in Table including compute overhead, minimum memory budget, search latency, memory fragmentation rate, and cutoff memory budget. For all of the five metrics, lower better. Compute overhead is the additional computation time due to tensor recomputation. Minimum memory budget is the minimum memory that the neural network can be trained on. Note that it does not include the memory of CUDA context. Search latency is the time to search for the optimal tensors to evict until the freed memory is sufficient for allocating the next tensor. Memory fragmentation rate is defined as the size free memory divided by the memory held by the memory allocator. Cutoff memory budget is the memory threshold, below which at least one tensor is evicted during the training process. The experimental results of the cutoff memory budget are displayed in Appendix C.
4.3 Results and Discussion
Compute Overhead and Memory Budget. Figure displayss memory budgets and the lowest memory budget for all methods. Coop presents the lowest compute overhead and can be trained under the minimum memory budgets on all eight networks. This improvement is substantial in most networks. For instance, the overhead of training BERT Large
displays the compute overhead across various

---

1.5
Large
(5)
BILSTM (2048)
Swin-T (40)
144
Compute Overhoad (60
12
1.2
Style 2:78 (8)
ResNet-50 (115)
Incaption V3 (96)
SPOS (70)
1.20
1.10
1.00
0.5
1.5
0.4
0.5
0.0
1.0
Memory Ratio
Coop (Oun)
Figure 4: Comparison among different tensor rematerialization strategies on eight DNNs, showing the ratio of compute overhead under different memory budgets (fractions of the original peak memory usage). Model batch sizes are given in parentheses. SAR (selective activation recomputation) is only available transformer models (BERT Large and GPT-3).
under 50% memory ratio with Coop is ,, compared with 29% with DTE and 4 with DTR. Additionally, Coop
up to 60% of memory usage when training Inception V3, while both DTR and DTE only save half of that amount.
We observed that the runtime overhead of training ResNet-50, Inception V3, and Swin-T is low for all three strategies under high memory ratios. However, the overheads of training the other DNNs, such as U-Net, BiLSTM, and BERT Large, are substantially higher with DTR and DTE than with Coop even at 90% memory ratio. The reason is that the memory size of the intermediate feature is decreasing in ResNet-50, Inception V3, and Swin-T, which is typical for image classification tasks. Evicting an old single tensor with a large size is often sufficient for allocating a new tensor with small size. By contrast, BERT Large and STM have features with consistent sizes and U-Net has a unique U-shaped architecture. In these cases, DTR and DTE need to run loops to release enough memory, which hugely increases the runtime overhead.
Search Latency. As shown in Figure the search latency of Coop is less than DTR and DTE in most cases. Moreover, the search latency of DTR and DTE displays considerable variations across the eight DNNs. The maximum search latency for training BiLSTM exceeds 10 ps using DTR. which is five orders of magnitude slower than the 0.32 for training ResNet-50. In contrast, Coop yields more consistent search latency across different DNNs. This because Coop only requires a single search using the efficient sliding window algorithm, and thus the latency remains low and is affected by the fluctuations the size of the intermediate tensors.
BERT Large
U-Net
BiLSTM
Swin-T
10a
5.0
10a
1.0
Search Latency (us)
10
10
0.5
10
0.0
0.0
0.5
0.6
0.7
0.8
0.4
0.5
0.6
0.7
0.8
0.5
0.6
0.7
0.8 0.9
0.3 0.6 0.8
GPT-3
Style
2.78
ResNet-50
Inception V3
SPOS
10
0.3
20
0.2
10
1 m
0.6
0.7
0.8
0.9
0.8
0.4 0.7 0.8
0.4
0.5
0.6
0.7
0.8
Memory Ratio
Coop (Qurs)
Figure 5: Search latency of Coop. DTR, and DTE under different memory ratios and networks. The bars with slashes represent the OOM error. The y-axes of GPT-2, BERT Large, and BiLSTM are in logarithmic scale.

---

We observed that Coop's search latency is larger than DTR and DTE on ResNet-50 under high memory ratios. We surmise this the result of different engineering implementations rather than the strategies. During the training of ResNet-50, evicting a single resident tensor is sufficient for allocating new one. Therefore, the search latencies of ResNet-50 are already very low (less than 0.4 ju)), and thus the influence of the implementation is more obvious.
Memory Fragmentation. As shown in Figure 6 training the eight DNNs using DTR leads to considerable memory fragmentation rates. Although DTE mitigates memory fragmentation compared DTR in most cases, its highest memory fragmentation rate remains considerably high. In contrast, Coop significantly reduces the memory fragmentation rate to less than 5 % across all DNNs and memory budgets. The performance of Coop is substantially better than DTR and DTE, particularly for BiLSTM and BERT Large models. This demonstrates that Coop utilizes the memory resources more efficiently and explains why the same DNNs can be trained with less compute overhead and lower memory budgets with Coop (Figure
BERT Large
U-Net
BiLSTM
Swin-T
10a
10
Avg Mem Frag (%4
20
5
10
0
0
0.5
0.6
0.7
0.8
0.9
0.4
0.5
0.6
0.7
0.8 0.9
0.5
0.6
0.7
0.8 0.9
0.3
0.4 D.S 0.6 0.7 0.8 0.9
GPT-3
Style
2.7B
ResNet-50
Inception V3
SPOS
10
10
5
10
5
5
0.6
0.7
0.8
0.9
0.3 0.4 0.5 0.6 0.7
0.8
0.9
0.4 0.5 0.6
0.7
0.8
0.9
0.4
0.5
0.6
0.7
0.8
0.9
Memory Ratio
Coop (Ours)
DTE
DTR
Figure 6: Averaged memory fragmentation rate under Coop, DTR, and DTE. lower rate denotes more efficient utilization of the runtime memory. The bars with slashes represent the OOM error.
5
Limitation
Coop comes with its own memory pool, so it cannot be simultancously used with CUDA's built-in memory pool (stream-ordered memory allocator). The advantage of using the stream-ordered memory allocator that multiple programs that use the stream-ordered memory allocator can share the same memory pool. However, all existing deep learning frameworks use their own memory pools instead of CUDA's built-in memory pool in default to achieve better efficiency and flexibility.
Additionally, Coop an online method so it inherits the pros and cons of online methods. It can provide efficient solution to finding the best tensors to evict within negligible time. However, the solutions may not match the optimal solutions as solved by the offline methods such as Checkmate, even though these offline methods usually require additional solvers and take several hours or multiple days.
Conclusion
We proposed Coop address the challenge training large dynamic DNNs under restricted memory budgets. We argued for the first time that the memory system in DL frameworks should not be ignored in tensor rematerialization and the conventional checkpointing algorithms on DCGs can be hugely improved by co-optimizing tensor allocation and tensor rematerialization. Leveraging the contributions the sliding window algorithm, cheap tensor partitioning, and recomputable in-place, Coop demonstrates consistent and superior performance over the state-of-the-ar checkpointing algorithms on most of the commonly used models, with the lowest memory budgets, least runtime overhead, and minimum memory fragmentation rate.
10

---

Acknowledgement
The authors thank Marisa Kirisame for discussing with us the idea and paper-writing of Coop. and thank Kan Wu, Qiaoling Chen, and Yipeng Li for their helpful advice. Jianhao Zhang would like to express heartfelt appreciation to his wife, Xingzi Yao, for her unwavering support and companionship. This work was supported by the Major Scientific Research Project of Zhejiang Lab 200DDDDO))) and National Natural Science Foundation of China under Grant No. U20A20226.
References
[1] Tom Brown, Benjamin Mann, Nick Ryder, Mclanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[2] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ily Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pages 8821-8831. PMLR, 2021.
[3] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anm 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.
[4] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Bureu Karagol Ayan, Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205. 1487, 2022.
[5] Amir Gholami, Zhewei Yao, Sehoon Kim, Michael W. Mahoney, and Kurt Keutzer. Ai and memory wall, Mar 2021.
[6] Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. Zero- infinity. Proceedings of the International Conference for High Performance Computing, Net- working. Storage and Analysis, Nov 2021.
[7] Andreas Griewank. An implementation of checkpointing for the reverse or adjoint model of differentiation. ACM Trans. Math. Software, 26(1):1-19 1999.
[8] Andreas Griewank and Andrea Walther. Algorithm 799: revolve: an implementation of check- pointing for the reverse or adjoint mode of computational differentiation. ACM Transactions on Mathematical Software (TOMS), 26(1):19 45, 2000.
[9] Atilim Gunes Baydin, Barak A Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: a survey. Journal of Marchine Learning Research, 18:1-43, 2018
[10] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint 2xivv:604:, 06174, 2016.
[11] 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 Muchine Learning and Systems 2020, pages 497-511. 2020.
[12] Ravi Kumar, Manish Purohit, Zoya Svitkina, Erik Vee, and Joshua Wang. Efficient rema- terialization for deep networks. Advances in Neural Information Processing Systems, 32, 2019.
[13] Audrunas Gruslys, Rémi Munos, Ivo Danihelka, Marc Lanctot, and Alex Graves. Memory- efficient backpropagation through time. Advances Neural Information Processing Systems, 29, 2016.

---

[14] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vito,, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8024-8035. Curran Associates, Inc., 2019.
[15] Vijay Korthikanti, Jared Casper. Sangkug Lym, Lawrence McAfee, Michael Andersch, Moham- mad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. arXiv preprint arXi::2005.05/96 2022.
[16] Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, and Zachary Tatlock. Dynamic tensor rematerialization. arXiv preprint arXiv:2006.00 09666, 2020.
[17] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp. Geoffrey Irving, Michael Isard, Rafal Jozefowicz, Yangqing Jia, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mané, Mike Schuster, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Jonathon Shlens, Benoit Steiner, llya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow, Large-scale machine learning on heterogeneous systems, 11 2015.
[18) Tianqi Chen, Mu Li, Yutian Li, Min Lin, Naiyan Wang, Minjie Wang, Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang. Mxnet: A flexible and efficient machine learning library for heterogeneous distributed systems. arXiv preprint arXiv:1512.01274 2015.
[19] Daan Leijen, Benjamin Zorn and Leonardo de Moura. Mimalloc: Free list sharding in action.
In Asian Symposium on Programming Languages and Systems, pages 244-265. Springer, 2019.
[20] Jason Evans. A scalable concurrent malloc (3) implementation for freebsd. In Proc. of the bsdcaan conference, ollawa, canada, 2006.
[21] Martin Maas, David G. Andersen, Michael Isard, Mohammad Mahdi Javanmard, Kathryn McKinley, and Colin Raffel. Learning- based memory allocation for c++ server workloads. In 25th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2020.
[22] Chang Hyun Park, Sanghoon Cha, Bokyeong Kim, Youngjin Kwon, David Black-Schafft and Jaehyuk Huh. Perforated page: Supporting fragmented memory allocation for large pages, In Proceedings of the ACM/EEEE 47th Annual International Symposium on Computer Architecture, ISCA '20, page 913-925. IEEE Press, 2020.
[23] Zhongzhe Hu, Junmin Xiao, Zheye Deng, Mingyi Li, Kewei Zhang, Xiaoyang Zhang, Ke Meng, Ninghui Sun, and Guangming Tan. Megtaichi: dynamic tensor-basee memory management optimization for dnn training. In Proceedings of the 36th ACM International Conference on Supercomputing, pages 1-13, 2022.
[24] Yu Tang, Chenyu Wang, Yufan Zhang, Yuliang Liu, Xingcheng Zhang, Linbo Qiao, Zhiquan Lai, and Dongsheng Li. Delta: Dynamically optimizing gpu memory beyond tensor recomputation, 2022.
[25] Parrots Team. Parrots framework. httpa///pprrotsddc.reaatttteeee
overview. html 2015.
[26] Daan Leijen. Koka: Programming with row polymorphic effect types. Electronic Proceedings in Theoretical Computer Science, 153:100-126, Jun 2014.
[27] Alex Reinking, Ningning Xie, Leonardo de Moura, and Daan Leijen. Perceus: Garbage free reference counting with reuse. In Proceedings of the 42nd ACM SIGPLAN International Conference Programming Language Design and Implementation, pages 96-111, 2021.
12

---

[28] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im Im: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv: (95V::/900000055: 2019.
[29] Aidan N Gomez, Mengye Ren, Raquel Urtasun, and Roger B Grosse. ual network: Backpropagation without storing activations. Advances in neural information processing systems, 30, 2017.
The reversible resid-
[30] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprini arXiv:2001.044511 2020.
[31] Animesh Jain, Amar Phanishayee, Jason Mars, Lingjia Tang, and Gennady Pekhimenko. Gist; Efficient data encoding for deep neural network training. In 2018 ACM/IEEE 45th Annual International Symposium on Computer Architecture (ISCA), pages 776-789. IEEE, 2018.
[32] Chien-Chin Huang, Gu Jin, and Jinyang Li. Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping. In Proceedings of the Twenry-Fifhh International Conference on Architectural Support for Programming Languages and Operating Systems, pages 1341-1355, 2020.
[33] Olivier Beaumont, Lionel Eyraud-Dubois and Alena SHILOVA. Efficient combination of rematerialization and offloading for training DNNs. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[34] Zhihao Jia, Matei Zaharia, and Alex Aiken. Beyond data and model parallelism for deep neural networks. Proceedings of Machine Learning and Systems, 1:1-13, 2019.
[35] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pages 3-19, 2018.
[36] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. International conference on machine learning, pages 448 456. pmlr, 2015.
[37] Jianfei Chen, Lianmin Zheng, Zhewei Yao, Dequan Wang, Ion Stoica, Michael W Mahoney, and Joseph E Gonzalez. Acton: Reducing training memory footprint via 2-bit activation compressed training. In International Conference on Machine Learning, 2021.
[38] Vijay Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Moham- mad Shocybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. arXiv preprint arXiv:2055.05/99 2022.
[39] Jinhui Yuan, Xinqi Li, Cheng Cheng, Juncheng Liu, Ran Guo, Shenghang Cai, Chi Yao, Fei Yang, Xiaodong Yi, Chuan Wu, Haoran Zhang, and Jie Zhao. Oneflow: Redesign the distributed deep learning framework from scratch, 2021.
[40] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory opti- mizations toward training trillion parameter models. SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, Noy 2020.
13


Reviews:
[{"review_id": "QsZmfuK5ab", "reviewer": "Reviewer_BV8L", "summary": "This paper considers the rematerialization problem for DNN training and studies it from the perspective of memory fragmentation.", "questions": "- I think this paper is interesting in the sense that it raises and studies a problem that could affect the performance of other rematerialization algorithms in the literature. The problem of memory fragmentation is not taken into account in most of the DNN memory optimization papers in the literature.\n\n- The comparisons presented in the paper seem to be limited to heuristic based methods only. I think a comparison against the checkmate method ([11]) would make the results more interesting. The checkmate method is known to return the optimal solution since it's an exact method. However, it's also known that it doesn't scale to large-scale graphs. Perhaps numerical experiments for some small scale graphs could be still valuable since checkmate as a baseline would represent the optimal under the assumption that memory fragmentation is not an issue. This, in my opinion, would make the contributions of this paper clearer.\n\n- I haven't read the work of [21]. Given what the authors discuss about the DTE method, is this statement in line 64 true: \"We argued for the first time that existing tensor rematerialization methods overlook the memory system during optimization and wrongly assume that the memory in DL systems is a fungible commodity\"? Another statement similar to this is \"To the best of our knowledge, Coop is the only tensor rematerialization scheme that fully bypasses the incorrect assumption of DL memory system.\" \n\nMinor\n- It is not immediately clear what is meant by \"search latency\" the first time it is mentioned in the text. It is explained later in the text, perhaps that explanation could be moved up to where it's mentioned first in the text.\n- This sentence in line 70 is hard to follow: \"The properties of memory allocators in deep learning frameworks are considered to reduce the heuristic ...\" ", "limitations": "Please see the \"Questions\" section.", "rating": 6, "confidence": 3, "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 thank the reviewer’s appreciation for our work. Please see below our responses to your comments.\n\n> **Q2:** Comparison between Coop and Checkmate\n\nWe agree that a comparison with Checkmate could further demonstrate Coop's contributions. Checkmate's publicly available code includes two git branches, namely mlsys20_artifact and master.\n\nThe mlsys20_artifact branch is designated for replicating experiments from their paper. As the reviewer mentioned, these experiments do not account for memory fragmentation. What's more, Checkmate in this branch cannot generate executable networks (as evident in the tests/test_execution.py file within the mlsys20_artifact branch). Consequently, it cannot be used to investigate the impact of memory fragmentation.\n\nThe master branch is capable of generating executable TensorFlow2 Keras computational graphs. Running these actual Keras graphs should help us comprehend memory fragmentation's effects on Checkmate. However, in the case of ResNet50, we ran Checkmate for 14 hours without obtaining any results within a single designated memory budget, despite leveraging the advanced MILP solver, Gurobi. Furthermore, attempts to optimize U-Net using Checkmate yielded nearly identical networks regardless of the budget value specified, rendering the collected data nonsensical.\n\nWe will continue to make ongoing efforts, aiming to include a Coop and Checkmate comparison in the revised manuscript's appendix. Additionally, reviewers can refer to the comparison experiments between DTR and Checkmate in [1]. These experiments also do not account for memory fragmentation and do not generate executable networks. They reveal that DTR and Checkmate exhibit comparable performance across three networks (VGG-16, MobileNet, U-Net). In our experiments considering memory fragmentation, Coop outperforms DTR significantly (Figure 4), and DTR displays pronounced memory fragmentation issues (Figure 5). We believe this can serve as a supplementary rough reference.\n\n> **Q3:** Is Coop the only tensor rematerialization scheme that fully bypasses the incorrect assumption of DL memory system?\n\nThe heuristic in DTE encourages evictions of tensors that are adjacent to free memory blocks. However, DTE is still a greedy algorithm that runs in a loop to find the best tensor to evict until the next tensor can be successfully allocated. This brings in the redundant and discontinuous evictions, as a limitation of assuming the memory in DL systems is a fungible conmmodity. In comparison, Coop optimizes 'tensor allocation' and uses the sliding algorithm to produce a continuous block. Therefore, we claimed that Coop is the first to fully bypass this incorrect assumption.\n\n**Minor points:**\n\n> **1.** It is not immediately clear what is meant by \"search latency\" the first time it is mentioned in the text. It is explained later in the text, perhaps that explanation could be moved up to where it's mentioned first in the text.\n\nWe will define search latency at the position of its first occurrence.\n\n> **2.** This sentence in line 70 is hard to follow: \"The properties of memory allocators in deep learning frameworks are considered to reduce the heuristic ...\"\n\nWe will rephrase this sentence as: The heuristic of tensor rematerialization is reduced by taking into account the properties of memory allocators, and the memory allocators are improved by considering the efficiency of different operations in tensor rematerialization.\n\n[1] Kirisame M, Lyubomirsky S, Haan A, et al. Dynamic tensor rematerialization[J]. arXiv preprint arXiv:2006.09616, 2020."}, {"review_id": "qWS1zEy98j", "reviewer": "Reviewer_bdP2", "summary": "This paper proposes an optimization framework called Coop to solve the severe memory fragmentation, which is overlooked by prior tensor rematerialization works. Coop designs a sliding window algorithm to determine evicted tensor, guaranteeing the freed memory is contiguous and available for a new tensor.  Further, Coop adopts a cheap tensor partitioning method to rearrange the tensor in the memory layout based on the cost density, and a memory reuse mechanism, namely recomputable in-place, for the in-place operations. There two ideas are combined to reduce additional tensor rematerializaiton cost.", "questions": "1. From the figures in the evaluation part, the proposed Coop is not always optimal in search latency. Please analyze the underlying reasons.\n2. Is the additional sliding window search algorithm needed after recomputable in-place and cheap tensor partitioning in Figure 1? If not, how can the evicted tensors be determined with minimum cost? It would be better to provide an algorithm description. ", "limitations": "1. Is there any trade-off for implementing the framework Coop? ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper provides a framework to solve the bottleneck of the high memory fragment rate in the tensor rematerialization scheme. The proposed sliding algorithm, cheap tensor partition mechanism, and recomputable in-place method reduce the rematerialization cost and improve memory utilization. \n2. The problem formulation and writing make the paper is easy to understand.", "weaknesses": "1. It would be better to give an algorithm to describe the framework comprehensively.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer’s appreciation for our work. Please see below our responses to your comments. We use Q, W, L to denote question, weakness, and limitation correspondingly.\n\n> **W1:** It would be better to give an algorithm to describe the framework comprehensively.\n\n   We greatly appreciate your valuable comments. The pseudo-code algorithm for Coop is displayed below, with the core logic within the `allocate` function. The implementations of methods such as `free_and_merge_block` are omitted since they are straightforward and not directly related to the core logic. The pseudo-code will be added to Section 3.2 of the revised manuscript.\n\n```python\ndef evict(tensor):\n    # Release memory block and merge if possible\n    free_and_merge_block(tensor.addr)\n    tensor.addr = None\n\ndef rematerialize(tensor):\n    if tensor.addr is not None:\n        return\n    for x in tensor.producer_op.inputs:\n        rematerialize(x)\n    run(tensor.producer_op)\n\ndef run(op):\n    output_size = infer_size(op)\n    output = allocate(output_size, op)\n    # ...\n    if is_inplace_mutation(op):\n        op.inputs[0].addr = None\n\ndef allocate(size, producer_op):\n    if is_inplace_mutation(producer_op):\n        # Apply recomputable in-place\n        # For simplicity, we assume in-place mutation operations have only one input tensor.\n        # The addr of the input tensor will be set to None at the end of `run` method.\n        addr = producer_op.inputs[0].addr\n    else:\n        block = find_free_block_large_than(size)\n        if block is None:\n            # Apply sliding window algorithm\n            evict(sliding_window_search(size))\n            block = find_free_block_large_than(size)\n        # Apply cheap tensor partitioning\n        if is_expensive(op_type):\n            addr = block.left_addr\n        else:\n            addr = block.right_addr - size\n    return Tensor(addr, size, producer_op)\n```\n\n> **Q1:** Why is Coop not always optimal in search latency?\n\n   We provided an explanation in Line 310. Given the structure of ResNet-50, evicting a single resident tensor might be sufficient to allocate a new one. In this case, both DTR and DTE could find the best tensors to evict after one iteration. Therefore, the search latencies depend on the engineering implementation rather than the strategies themselves.\n\n> **Q2:** Is the additional sliding window search algorithm needed after recomputable in-place and cheap tensor partitioning in Figure 1? If not, how can the evicted tensors be determined with minimum cost?\n\n   The reviewer's understanding is correct that the memory layout is first optimized during tensor allocation by using recomputable in-place and cheap tensor partitioning. Sliding window algorithm is used to find the best tensors to evict given the current memory layout. The three modules are flexibly used during the whole training process. The sequence of the three modules in Figure 1 is not the sequence of using them in time. We will modify the caption of Figure 1 and add the pseudo-code in the answer of W1 to the revised manuscript to make the whole pipeline clearer.\n\n> **L1:** Is there any trade-off for implementing the framework Coop?\n\n   Coop comes with its own memory pool, so it cannot be simultaneously used with CUDA's built-in memory pool (stream-ordered memory allocator). The advantage of using the stream-ordered memory allocator is that multiple programs that use the stream-ordered memory allocator can share the same memory pool. However, all existing deep learning frameworks use their own memory pools instead of CUDA's built-in memory pool in default to achieve better efficiency and flexibility.\n\n   Additionally, Coop is an online method so it inherits the pros and cons of online methods. It can provide an efficient solution to finding the best tensors to evict within negligible time. However, the solutions may not match the optimal solutions as solved by the offline methods such as Checkmate, even though these offline methods usually require additional solvers and take several hours or multiple days.\n\n   We will add these discussions about limitations to the revised manuscript."}, {"review_id": "SH9GioSqNf", "reviewer": "Reviewer_ks3C", "summary": "Tensor materialization trades the memory with recomputation. Prior tensor materialization methods do not consider the memory fragmentation problem of the memory system used in deep learning frameworks, which makes them evict unnecesary tensors. The authors of this paper proposed a memory-system-aware rematerialization method called Coop to reduce the memory fragmentation. Experiments show that Coop can achieve up to 2x memory saving comparied with prior works. ", "questions": "1. Whether the prior works have the fragmentation problem if they use the memory pool implemented in cuda? (See weaknesses part)\n", "limitations": "The proposed method relied on that the underlying memory allocator is able to merge contiguous chunks. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The authors proposed a new method based on sliding window algorithm (sec 3.3) to alliviate the fragmentation problem of rematerialization.\n2. The experiments show that this method is effective and can greatly reduce the compute overhead for specific memory ratio compared with prior works.\n\n", "weaknesses": "1. The requirement of underlying memory allocator might limit the applicability of the proposed method\n\nThis paper assumed the memory allocator used by the deep learning systems (discussed in section 2.1). The underlying memory allocator must be able to **merge** the freed chuncks if they are contiguous. There are other kinds of memory allocators (e.g., record the mapping from chunck size to a list of free chuncks with the specific chunk size) that do not have this feature, thus the proposed method can not be (directly) used for deep learning systems with this kind of memory allocators.\n\n2. More discussion on the effectiveness of the the page-table-based memory system is needed\n\nSimilar to CPU memory system, the GPU memory system also employed the page table to manage its memory. Thus, we can free the discontiguous memory chunks and allocate a new one with the sum of the sizes of the freed chunks. From the virtual memory's view, the allocated memory is contiguous. Thus, there is no fragmentation problem discussed in this paper, and the prior works can be directly used.   The good news is that, since CUDA 11.2, we can directly use the memory pool [1] implemented in cuda runtime/driver to enjoy this feature. Thus, I am interested in whether the prior works have the fragmentation problem if they use the memory pool implemented in cuda?\n\nWe can use the following program to validate that the page-table based GPU memory system. In the program, we allocate 3 chuncks of 7 GB called p1, p2, and p3. Then, we free p1 and p3 (they are not contiguous). At last, we allocate a chunk of memory with 14 GB called p4 successfully. \n```python\n#include <stdio.h>\n#include <cuda.h>\n#include <unistd.h>\n\n#define CHECK(e) {auto s = e; if (s != cudaSuccess) printf(\"CUDA error: %s\", cudaGetErrorString(s));}\n#define GB(x) ((x) * 1024ull * 1024 * 1024)\n\nint main() {\n    void *p1, *p2, *p3, *p4;\n    CHECK(cudaMalloc(&p1, GB(7)));\n    CHECK(cudaMalloc(&p2, GB(7)));\n    CHECK(cudaMalloc(&p3, GB(7)));\n    printf(\"first allocation done\\n\");\n    printf(\"p1=%p\\n\", p1);\n    printf(\"p2=%p\\n\", p2);\n    printf(\"p3=%p\\n\", p3);\n\n    // sleep 5 seconds, we can check the memory usage in `nvidia-smi` during this time\n    sleep(5);\n\n    printf(\"free p1 and p3, allocate p4\\n\");\n    CHECK(cudaFree(p1));\n    CHECK(cudaFree(p3));\n    CHECK(cudaMalloc(&p4, GB(14)));\n    printf(\"p4=%p\\n\", p4);\n\n    printf(\"successfully allocated p4\\n\");\n    CHECK(cudaFree(p2));\n    CHECK(cudaFree(p4));\n}\n```\nI can run above program on a RTX 3090 (24GB memory).\n```\nfirst allocation done\np1=0x7f4a68000000\np2=0x7f48a8000000\np3=0x7f46e8000000\nfree p1 and p3, allocate p4\np4=0x7f4528000000\nsuccessfully allocated p4\n```\n\n[1] https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY__POOLS.html", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer’s appreciation for our work. Please see below our responses to your comments. We use Q, W, L to denote question, weakness, and limitation correspondingly.\n\n> **W1&L1:** The underlying memory allocator must be able to merge the freed chunks if they are contiguous.\n\nThe reviewer's understanding is correct that Coop requires the memory allocator to be able to merge the freed chunk if they are contiguous. All DL frameworks have implemented 'merge', e.g., `try_merge_blocks` method in PyTorch's caching allocator and `Merge` method in TensorFlow's BFC allocator. Although 'merge' can only be used to combine the sub-block in the same block, it is common to apply for a single block in the size of the memory budget when rematerialization is enabled (so every sub-block is from the same block). Therefore, we suppose this limitation will not influence the wide application of Coop in most DL frameworks. We will add these discussions in the revised manuscript.\n\n> **W2&Q1:** More discussion on the page-table-based memory system is needed.\n   \n   We are deeply appreciative of the reviewer's insightful comments. We modified the attached codes to use the memory pool implemented in CUDA:\n   \n   ```cpp\n   #include <stdio.h>\n   #include <cuda_runtime_api.h>\n   #include <cuda.h>\n   #include <unistd.h>\n   #include <cassert>\n   \n   #define CHECK(e) {auto s = e; if (s != cudaSuccess) printf(\"CUDA error: %s\", cudaGetErrorString(s));}\n   #define GB(x) ((x) * 1024ull * 1024 * 1024)\n   \n   int main() {\n       {\n           // assert current device supported memory pool\n           int value = 0;\n           CHECK(cudaDeviceGetAttribute(&value, cudaDevAttrMemoryPoolsSupported, 0));\n           assert(value == 1);\n       }\n       void *p1, *p2, *p3, *p4;\n       cudaStream_t stream;\n       CHECK(cudaStreamCreate(&stream));\n   \n       CHECK(cudaMallocAsync(&p1, GB(2), stream));\n       CHECK(cudaMallocAsync(&p2, GB(2), stream));\n       CHECK(cudaMallocAsync(&p3, GB(2), stream));\n       printf(\"first allocation done\\n\");\n       printf(\"p1=%p\\n\", p1);\n       printf(\"p2=%p\\n\", p2);\n       printf(\"p3=%p\\n\", p3);\n   \n       printf(\"free p1 and p3\\n\");\n       CHECK(cudaFreeAsync(p1, stream));\n       CHECK(cudaFreeAsync(p3, stream));\n       // Sleep 1 seconds to imitate the time between free and alloc.\n       // We can sleep arbitrary time here and the result is the same.\n       sleep(1);\n       printf(\"allocate p4\\n\");\n       CHECK(cudaMallocAsync(&p4, GB(2.1), stream));\n       printf(\"p4=%p\\n\", p4);\n   \n       if (p4 != nullptr) {\n         printf(\"successfully allocated p4\\n\");\n         CHECK(cudaFreeAsync(p4, stream));\n       } else {\n         printf(\"failed to allocate p4\\n\");\n       }\n       CHECK(cudaFreeAsync(p2, stream));\n       CHECK(cudaStreamDestroy(stream));\n   }\n   ```\n   \n   We ran the program on an NVIDIA GeForce RTX 2080 with 8GB memory and the chunk p4 **cannot** be successfully allocated. This indicates that the memory pool provided by CUDA did not fully utilize page table to optimize the memory usage. We further do three more experiments:\n   \n   * **Experiment 1:** Replace `sleep(1)` with `cudaStreamSynchronize(stream)`. According to the CUDA documentation, `cudaStreamSynchronize` returns all free memory blocks in the CUDA memory pool to OS. \n      \n      **Result:** The allocation of `p4` succeeds.\n   * **Experiment 2:** Free `p1` and `p2` (or `p2` and `p3`) instead of `p1` and `p3`, so that the freed memory is contiguous in the virtual memory space.\n      \n      **Result:** The allocation of `p4` succeeds.\n   * **Experiment 3:** Based on Experiment 2, change the size of `p4` from 2.1GB to 4GB and 4.1GB, so that the desired memory size is exactly equal to or slightly larger than the freed memory size.\n      \n      **Result:** The allocation of `p4` succeeds and fails, respectively.\n   \n   These three experiments suggest that the memory pool provided by CUDA also caches the freed memory block, just like what the normal allocator of deep learning frameworks does. Based on these experimental results, we believe the prior works have the fragmentation problem even if they use the memory pool implemented in CUDA.\n   \n   However, we fully agree that optimizing memory allocation by combining page table to reduce memory fragmentation is a promising direction. Studies have proven that memory fragmentation in CPU can be reduced by using similar ideas [1,2]. We believe we can apply the same idea to optimizing the general memory allocation (not only recomputation) of deep learning systems if we can manipulate the underlying operations of the GPU driver (rather than treating it as a proprietary NVIDIA-controlled black box). We will add these discussions to the revised manuscript.\n   \n   [1] Maas M, Andersen D G, Isard M, et al. Learning-based memory allocation for C++ server workloads[C]//Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems. 2020: 541-556.\n   \n   [2] Park C H, Cha S, Kim B, et al. Perforated page: Supporting fragmented memory allocation for large pages[C]//2020 ACM/IEEE 47th Annual International Symposium on Computer Architecture (ISCA). IEEE, 2020: 913-925."}, {"review_id": "WE5BnTil3Q", "reviewer": "Reviewer_UBH8", "summary": "The authors propose to consider memory fragmentation when using tensor rematerialization on dynamic computation graphs. With three new memory management methods, sliding window algorithm, cheap tensor partitioning, and recomputable in-place, the authors enhance the checkpointing method with lower computation overhead and memory consumption.\n", "questions": "Please see the Weaknesses, especially the major issues.", "limitations": "The authors do not discuss limitations. I think the proposed COOP has the inherent limitations of the checkpoint method on the dynamic computation graphs.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is well-organized and easy to follow.\n* The perspective is new. Considering memory fragmentation can enhance the gradient checkpointing.\n* The results are great across the eight neural networks.\n* The implementation is available in the supplementary material.\n", "weaknesses": "**Major issues**\n1. Table 1. The cost density depends on the hardware accelerators. It is better to introduce the hardware specifications. The authors can also mention the arithmetic intensity, which is independent of the hardware. Also, in a single neural network, different operators belonging to the same category may have different cost densities (e.g., different kernel sizes in convolution layers).\n2. In cheap tensor partitioning, the authors allocate tensors from the leftmost and rightmost of the memory pool. Is it possible to have more ports for the memory pool? What if there are 3 levels of cost densities in the computation graph? In general computation graphs, the cost density or the arithmetic intensity of operators has a continuous distribution. Is it a good idea to classify them into two categories, “cheap” and “expensive” tensors?\n3. The authors miss one of the most important evaluation metrics, the end-to-end training time under different memory budgets, which is more critical than the computation overhead.\n4. The authors mention that “static graph methods are beyond the scope of this paper.” I am aware that symbolic and eager executions have distinct features. However, I have several concerns about that.\n    * Among the eight computation graphs in the experiments, six have static structures, while only two have dynamic ones. It is better to discuss more on the dynamic structures.\n    * The proposed method leverages the global information that several tensors are unevictable. If the computation graph is fully dynamic, we have no idea which tensors will be used in the future.\n\n**Minor issues**\n1. Line 180. Please explain the N. I assume that it is the number of tensors.\n2. A period after “Compute Overhead and Memory Budget” in Line 287, “Search Latency” in Line 303, and “Memory Fragmentation” in Line 315.\n3. The caption of Figure 5. y-axes -> y-axis.\n4. Results on Search Latency. The authors may list other statistics, e.g., min/max value, standard deviation to demonstrate, to demonstrate the small variations of the proposed method.\n5. The authors mainly use the term “memory allocation” in the paper. It is better to replace it with “memory management” since there are other memory operations, such as eviction.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer’s appreciation for our work, and we address the questions below. We use W and L to denote major weakness and limitation correspondingly.\n\n> **W1:** It is better to introduce the hardware specifications about cost density. The authors can also mention the arithmetic intensity. Different operators belonging to the same category may have different cost densities.\n   \n   We agree that the cost density depends on the hardware accelerators. The cost density in Table 1 is measured on NVIDIA GeForce RTX 2080 Ti. We will mention it in the revised manuscript.\n   \n   Arithmetic intensity is the ratio of FLOPs to the summation of input and output sizes. In the revised manuscript, we will define a similar concept, FLOPs density (the ratio of FLOPs to output size), to match the heuristic of recomputation, and will show FLOPs densities and cost densities of various operators (e.g., convolutions with different kernel sizes). Some preliminary results are shown in Table 1 of the attached pdf.\n   \n> **W2.1:** In general computation graphs, the cost density or the arithmetic intensity of operators has a continuous distribution. Is it a good idea to classify them into two categories, “cheap” and “expensive” tensors?\n   \n   We agree with the reviewer that cost density continuously changes. We divided the operations into two categories since we observed an obvious jump in the cost densities of the commonly used operations, as shown in Figure 1 of the attached pdf.\n   \n   Similar ideas of dividing the operations into two categories have been evidenced in Selective Activation Recomputation (SAR). SAR found that tensors generated by some operations (e.g., softmax and dropout) occupys 70% of the memory in GPT-3 while these operations only contribute to 2.7% FLOPs.\n\n> **W2.2:** Is it possible to have more ports for the memory pool?\n\n   It is possible to have more than two ports for the memory pool, where the operations are divided into two categories given the reasons in W2.1.\n\n   Figure 2(a)(1) shows an example of a 4-port memory pool with two 2-port sub-pools. The tensors are allocated to the two sub-pools alternatively. If there is no space left in one sub-pool, the tensors will be allocated to the other one. Figure 2(a)(2) shows an 8-port memory pool. Thus, we can construct memory pools with 16-port, 32-port, etc., in the same way.\n\n   If the ports are odd (2k+1), the memory pool is most likely to be consisted of k 2-port sub-pools and one 1-port sub-pool. This leads to insufficient usage of the memory as shown in Figure 2(a)(3).\n\n   The N-port memory pools can reduce the evictions of tensors that are continuous in the computational graph (thus reducing the heuristic) but will increase the memory fragmentation. We investigated this approach at the beginning of this work but did not apply it, since there is no significant improvement in the overall performance. Some experimental results of using 3-port and 4-port memory pools are shown in Figure 2(b). We will add these discussions to the revised manuscript.\n\n> **W3:** the end-to-end training time and the compute overhead.\n   \n   We followed Checkmate and DTR to use the metric of compute overhead, which can be regarded as the normalized training time (the end-to-end training time with rematerialization divided by the end-to-end training time without rematerialization).\n\n> **W4.1:** It is better to discuss more on the dynamic structures.\n   \n   We agree that online methods such as DTR and Coop are useful to optimize NNs with dynamic structures since offline optimization methods are not applicable to these dynamic NNs. We also studied multiple static NNs since (1) there are more static NNs than dynamic NNs, and (2) online methods also show advantages over offline methods when used to optimize static NNs. As in Reviewer BV8L’s comment, one popular offline method, Checkmate, is limited to running on small-scale networks. Our experiments for the Question 2 of Reviewer BV8L also show that Checkmate (with the most advanced commercial solver, gurobi) fails to find the best solution of ResNet-50 within a 14-hour time limit. This explains why PyTorch and TensorFlow only provide manual recomputation instead of the theoretically optimal method such as Checkmate. Given the small search latency of online methods such as DTR and Coop, we expect that applying online methods to optimize trainings of both dynamic or static NNs will be beneficial. We will add these discussions in the revised manuscript.\n\n> **W4.2:** Coop leverages the global information that several tensors are unevictable.\n   \n   The unevictable tensors in Coop are the parameters and buffers in DNN. These tensors are available from the initialization of the network (e.g., the constructor of nn.Module in PyTorch and tf.keras.Model in TensorFlow 2.0). Therefore, Coop only uses information available for optimization.\n   \n> **L1:** Coop has the inherent limitations of the checkpoint method on the dynamic computation graphs\n\n   We agree with the reviewer that Coop has some inherent limitations of online methods. We will add the discussions (as in the response to Reviewer 2’s L1) to the revised manuscript.\n\n**Minor issues**\n\n> **1.** The meaning of N at Line 180.\n\nThe meaning of N is explained at Line 157. We will explain it at Line 180 in the revised manuscript to make it clearer.\n\n> **2.** List min/max value, standard deviation of search latency.\n\nThanks for your valuable advice. We calculated these statistics (shown in the following table) and found that they are very useful to demonstrate the small variations of Coop.\n\n|   | Coop | DTE | DTR |\n| ------------- | ------------- | ---- | ---- |\n| Min  | 0.24  | 0.20 | 0.24 |\n| Max  | 18.6  | 13975.0 | 28019.0 |\n| Std  | 2.81  | 2349.1 | 4863.5 |\n\n> **3.** Missing periods, \"y-axes\" -> \"y-axis\" and \"memory allocation\" -> \"memory management\"\n\nWe will update the manuscript accordingly."}]

Author rebuttal:
We thank the reviewers' appreciation and valuable advice for our work. Some new figures and tables are in the attached pdf file.

Meta-review: The paper introduces "Coop", an optimization framework aimed at addressing the memory fragmentation challenge in tensor rematerialization for Deep Neural Network (DNN) training. This is achieved through a combination of sliding window algorithm, cheap tensor partitioning, and recomputable in-place mechanisms. The paper claims that Coop reduces memory costs and improves memory utilization.

**Strengths:**

- Addresses the overlooked issue of memory fragmentation in tensor rematerialization.
- Introduces a unique approach with the sliding window algorithm, tensor partitioning, and recomputable in-place method to improve memory management.
- Clear problem formulation and writing make the paper comprehensible.
- Empirical results demonstrate up to 2x memory savings compared to previous approaches.
- The manuscript is well-organized and clear, with supplementary implementation materials provided.

**Weaknesses:**

- Unclear terminologies at places, such as "search latency".
- The paper is often unclear in terms of hardware specifications, tensor partitioning, and end-to-end training time under different memory budgets.
- The paper does not delve deeply into the potential benefits and limitations of using the page-table-based GPU memory system.
- Other limitations need to be described more explicitly. For example: (a) it relies on a specific type of memory allocator that can merge contiguous chunks, potentially limiting its applicability to certain deep learning systems. (b) The focus seems more inclined towards static computation graphs. Any limitations around dynamic computation graphs need to be described.
- Limited comparisons in the paper. Notably, a comparison against the "checkmate" method is missing, which represents an optimal solution.

The paper offers valuable contributions and fresh perspectives on tensor rematerialization by addressing memory fragmentation. However, there are areas that require further clarification and explicit description of limitations. I strongly recommend the authors revise the paper to accommodate these weaknesses.
Decision: Accept (spotlight)