Efficient Combination of Rematerialization and Offloading for Training DNNs

1 Introduction . A general trend in machine learning , especially in NLP today , is to increase the size of models to improve accuracy . For example , NLP models like BERT [ 7 ] or GPT-3 [ 30 ] have billions of parameters and vision networks like EfficientNet [ 35 ] , AmoebaNet [ 34 ] or SENET [ 15 ] routinely have a hundred million parameters , which raises difficult memory issues on architectures like GPUs . Moreover , it was observed [ 29 ] that a sufficiently large batch-size is required to obtain good and fast convergence . In this paper we focus on the strategies for reducing the memory usage of activations that can be implemented when considering a single computation node , typically a single GPU . These techniques can of course be combined with the different possible parallelizations of the training phase , like data parallelism [ 39 , 27 , 11 ] , model parallelism [ 25 , 37 , 26 ] , filter or kernel [ 9 ] or image parallelisms [ 8 ] . The combination with parallel strategies is however out of the scope of the present paper and our goal here is to model , analyze the complexity and understand how to benefit simultaneously from rematerialization and offloading . Rematerialization [ 14 , 5 , 20 , 3 ] is a technique that consists in deleting from memory some of the activations computed during the forward phase as soon as they are no longer needed for the forward phase . As these activations are required again during the backward phase , they are then recomputed from other activations kept in memory . The idea is therefore to trade additional computations for lower memory requirements . The objective , in the context of rematerialization , is to compute a sequence of valid operations ( i.e . such that the inputs of any operation are actually in memory at the time of the operation ) consisting of elementary forward , backward and delete operations , and which has a minimum execution time among all sequences that satisfy a given memory constraint . Offloading [ 32 , 2 , 23 , 4 ] is a technique that has the same objective , i.e . to produce a sequence of operations satisfying a given memory constraint . However , in the case of offloading , each computation operation is executed exactly once.The idea here is therefore to trade transfers on the PCI bus between 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) the CPU and the GPU for lower memory requirements . The objective , in the context of offloading , is therefore to compute a sequence of valid operations consisting of elementary forward , backward , offload and prefetch operations , and which has a minimum execution time among all sequences that fulfill a given memory constraint . While rematerialization and offloading share the common goal of limiting memory consumption , they rely on different techniques and achieve this result by consuming independent resources , using either more computation or more transfers between a GPU and a CPU . In this paper we solve the problem of finding the optimal sequence combining rematerialization and activation offloading , which has not been addressed in the literature to the best of our knowledge . We discuss related work in Section 2 and model and notations in Section 3 . The algorithmic solution to combine optimally checkpointing and offloading is presented in Section 4 and experiments to assess the advantage of using both offloading and rematerialization , based on a large range of neural networks and architecture parameters , are presented in Section 5 . 2 Related work . Parallelism-based memory optimization A first approach to control memory consumption is to use parallelization . Two main sources of memory consumption in the training phase for deep neural networks are , on the one hand , the storage of the network weights and , on the other hand , the storage of activations . To limit the storage cost of activations ( for a fixed batch size ) , data parallelism [ 39 , 27 , 11 ] is the method of choice . Each batch is divided into mini-batches which are in turn processed in parallel on different computation resources . In this scheme , however , weights are replicated on each resource and collective communications are required to synchronize weights between different resources . When models become large , these collective communications can hinder the scalability of the data parallel approach . On the contrary , model parallelism [ 6 , 17 , 25 , 37 ] is used to distribute weights over different available resources . Weights of a neural network are then distributed on different computation resources , and communications happen between consecutive layers assigned to different resources . Nevertheless , this approach also requires sophisticated strategies [ 26 ] to save fewer activations and models in memory . The data can also be split along other dimensions . In [ 8 ] , large images are split into smaller images , the network is trained on these small images ( augmented by a halo ) and extra communications are needed in order to synchronize parameter updates . A similar strategy was proposed to use channel and filter parallelism in [ 9 ] . All these parallel strategies can be combined with the single node memory optimization strategies considered in the present paper . Rematerialization-based strategies Rematerialization techniques consist in keeping only a subset of activations in memory during the whole training phase , while the others are dynamically recomputed at runtime when they are needed for the backpropagation . This allows to explore a tradeoff between memory usage and computational cost . Rematerialization was first considered in the context of Automatic Differentiation ( AD ) [ 12 ] under the name `` checkpointing '' . In the context of AD , computational networks consist in very long and homogeneous chains , in which the forward activation corresponding to the i−th stage of the chain has to be kept into memory until the i−th backward stage . Rematerialization algorithms consist in determining in advance which forward checkpoints should be kept into memory and which ones should be recomputed from stored checkpoints when performing the backward phase . In the case of homogeneous chains closed form formulas providing the exact position of checkpoints were proposed in [ 13 ] . The use of rematerialization strategies was advocated for DNN in several papers [ 14 , 5 , 20 ] . An implementation of a simple periodic strategy [ 5 ] was provided in PyTorch [ 28 ] 1 , but it is only optimal for the restricted case of homogeneous chains and when activation discards are done only before the loss . However , DNN models are in general more complicated . More recently , this strategy has been extended to general heterogenous networks represented as sequential networks in [ 3 ] 2 . The case of non sequential networks was also addressed in the literature [ 19 , 21 , 18 , 22 , 10 ] , mostly under the simplifying assumption that discarding computed activations is only possible before the computation of the loss function [ 22 , 10 ] . In [ 19 ] the authors proposed an algorithm 1https : //pytorch.org/docs/stable/_modules/torch/utils/checkpoint.html 2https : //gitlab.inria.fr/hiepacs/rotor dealing with dynamic data flows , though based on an heuristical approach . The authors of [ 18 ] considered minimization of the runtime with the use of Integer Linear Programming . Even though this approach is suitable for arbitrary computational graphs with heterogeneous costs , solving this ILP is computationally very expensive even for intermediate depth networks . Offloading-based strategies Offloading [ 32 , 2 ] is an alternative approach which consists in offloading some of the forward activations from the memory of the GPU to the memory of the CPU , which is expected to be much larger . In [ 32 ] , the authors proposed a simple mechanism of memory virtualization , that introduces unnecessary synchronizations between data transfers and computations of later forward activations . This technique was improved in [ 2 ] . In both papers , the algorithmic strategies to decide which activations to offload into the main memory are relatively straightforward , and basically consist in sending to the CPU either all activations or only those of convolutional layers . In [ 4 ] , the complexity of the offloading problem was analyzed and several NP-completeness results were provided , depending whether activations can be offloaded partially or not . In order to reduce the overhead incurred by the communications , the use of compression was advocated in [ 33 ] . Memory virtualization was further considered in [ 24 , 23 , 16 , 38 ] . Finally , the speed-centric rematerialization from [ 5 ] enhanced with memory-centric rematerialization ( discards activations of every segment all the time ) was combined with the simple offloading approach from [ 32 ] in [ 36 ] . In addition , the authors in [ 31 ] also used rematerialization of [ 5 ] with a possibility of further offloading saved checkpoints to the CPU if rematerialization only is not enough to perform training under memory constraint . 3 Model and notations . 3.1 Forward , backward and offloading elementary operations . In what follows , our goal is to process the chain described by Figure 1 that corresponds to the processing of a mini-batch . This chain depicted by Figure 1 is defined by the set of operations F ` and B ` . The default forward operations F ` requires only the activation produced by the previous layer F ` −1 ( or the input mini-batch in the case of F1 ) , whereas the default backward operation B ` require the preceding gradient δ ` , the activation a ` −1 and the intermediate generated data ā ` necessary for the backward step . The distinction between a ` and ā ` is crucial in practice , since it allows to consider computation graphs which have a sequential structure , but are not purely sequential . In this setting , it is indeed possible for F ` to represent a complex operation ( any Direct Acyclic Graph of layers ) , and thus ā ` contains all the intermediate activations produced by F ` and necessary to compute B ` , whereas a ` only contains the output activation of F ` , that will be used by F ` +1 . In general , a ` can therefore be much smaller than ā ` . In order to manipulate different types of activations , we define in Table 1 a set of generic forward and backward operations , that can be directly translated in automatic differentiation packages such as tf.GradientTape for Tensorflow or torch.autograd for PyTorch . Operations are defined both by the type of operands they use and produce ( a ` or ā ` ) and by the treatment they make of the input data after processing ( deleted or kept in memory ) . Each operation is associated to a processing time and a memory overhead that takes place during the computation . Both the processing time and the memory overhead only depend on whether the operation is a forward or a backward task . In addition to the operations described in Table 1 , we also consider two operations O ` and P ` that correspond respectively to offloading to and prefetching from CPU memory . O ` and P ` can operate both on a ` −1 and ā ` −1 , depending on which type of activation has been produced . Thanks to CUDA streams , it is possible to asynchronously offload and prefetch activations from and to the CPU memory , while independent computations are performed .

Summary: The paper tackles the memory bottlenecks of by activations that limits scaling of model and batch sizes in DL training. The approach is to discover the optimal combination of activation re-materialization and offloading to CPU memory that will minimize the runtime overhead. The optimal solution is found using a proposed dynamic programming algorithm, called `pofo`, which determines an optimal sequence of re-materialization and offload decisions for the network layers, under the constraint of GPU memory and communication bandwidth to CPU. The paper also presents simulation-based results showing that `pofo` more efficiently reduces activation memory pressure compared to different SOTA approaches for different image models and different memory and communication bandwidth constraints.