SlowGuess/NeurIPS_2023_OCR huggingface dataset:
When dealing with operations that directly mutate the content of the input tensor, i.e., in-place operations, DTR introduces a 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 3 (c)). Allocating new memory for the copied tensor has two 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, if this tensor is unevictable, such as the parameter that is updated 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.
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-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-reuse 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 on 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 discontinuous 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 1(c)). In this way, the rate of memory fragmentation can be reduced to the least.
4.1 Evaluation Setup
We compared our proposed method, Coop, with two state-of-the-art DCG-based automatic rematerialization methods, DTR [16] and DTE (the rematerialization method introduced in MegTaichi [23]), across eight representative DNNs. The eight DNNs include a 2.7B parameter GPT-3 style large language model, Swin-Transformer, ResNet-50, etc. Additionally, we compared our method with selective activation recomputation [38] (the rematerialization method designed for Transformer models in Megatron-LM[28]) 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-optimized	Traversal Count
Coop	The proposed tensor rematerialization method.	✓	✓	✓	Single
DTE	Our impl. of Dynamic Tensor Evicting [23] in OneFlow.	✓	~*	X	Multiple
DTR	Our impl. of Dynamic Tensor Re-materialization [16] in OneFlow.	✓	X	X	Multiple
SAR	Our impl. of Selective Activation Recomputation [38] in OneFlow.	X	X	X	-
* MS is the abbreviation of the memory system.
* DTE's heuristic considers adjacent free memory but cannot promise a contiguous memory block is 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-aligned 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 1 to align with DTR and Coop. We also compared Coop in OneFlow with the official DTR in PyTorch and the official DTE in MegEngine in Appendix A. The effects of the three proposed modules were examined in Appendix B. All experiments were conducted on a machine equipped with 4 NVIDIA 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 metrics are used for evaluating the performance of the methods in Table 2, including compute overhead, minimum memory budget, search latency, memory fragmentation rate, and cutoff memory budget. For all of the five metrics, lower is 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 of 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 4 displays the compute overhead across various 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
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 on transformer models (BERT Large and GPT-3).
under a  $50\%$  memory ratio with Coop is  $11\%$ , compared with  $29\%$  with DTE and  $41\%$  with DTR. Additionally, Coop can save 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 a  $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 a small size. By contrast, BERT Large and BiLSTM have features with consistent sizes and U-Net has a unique U-shaped architecture. In these cases, DTR and DTE need to run in loops to release enough memory, which hugely increases the runtime overhead.
Search Latency. As shown in Figure 5, 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^{4}\mu s$  using DTR, which is five orders of magnitude slower than the  $0.32\mu s$  for training ResNet-50. In contrast, Coop yields more consistent search latency across different DNNs. This is because Coop only requires a single search using the efficient sliding window algorithm, and thus the latency remains low and is not affected by the fluctuations in the size of the intermediate tensors.
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 a 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 is 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 a new one. Therefore, the search latencies of ResNet-50 are already very low (less than 0.4  $\mu$ s), 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 to 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 4).
Figure 6: Averaged memory fragmentation rate under Coop, DTR, and DTE. A 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 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.
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.
We proposed Coop to address the challenge of 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 of the sliding window algorithm, cheap tensor partitioning, and recomputable in-place, Coop demonstrates consistent and superior performance over the state-of-the-art checkpointing algorithms on most of the commonly used models, with the lowest memory budgets, least runtime overhead, and minimum memory fragmentation rate.
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 (No.2019KD0AD01) and National Natural Science Foundation of China under Grant No. U20A20226.
[1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie 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 Ilya 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-Anne Lachaux, Timothé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 Sahara, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 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, Networking, 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 checkpointing 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 Machine 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 arXiv:1604.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. In Proceedings of Machine Learning and Systems 2020, pages 497-511. 2020.
[12] Ravi Kumar, Manish Purohit, Zoya Svitkina, Erik Vee, and Joshua Wang. Efficient rematerialization 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 in Neural Information Processing Systems, 29, 2016.
[28] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053, 2019.
[29] Aidan N Gomez, Mengye Ren, Raquel Urtasun, and Roger B Grosse. The reversible residual network: Backpropagation without storing activations. Advances in neural information processing systems, 30, 2017.
[30] Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451, 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 thegpu memory limit via smart swapping. In Proceedings of the Twenty-Fifth 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 Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In 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. Actnn: 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, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. arXiv preprint arXiv:2205.05198, 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, Olatunj Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, Nov 2020.