Efficient Combination of Rematerialization and Offloading for Training DNNs
Authors: Olivier Beaumont
Lionel Eyraud-Dubois
Alena SHILOVA
Venue: NeurIPS 2021 Poster (2021)
Keywords: combinatorial optimization algorithms
dynamic programming
rematerialization
offloading
checkpointing
memory constraint
training
deep neural networks
feed forward backpropagation training

Abstract:   Rematerialization and offloading are two well known strategies to save memory during the training phase of deep neural networks, allowing data scientists to consider larger models, batch sizes or higher resolution data. Rematerialization trades memory for computation time, whereas Offloading trades memory for data movements. As these two resources are independent, it is appealing to consider the simultaneous combination of both strategies to save even more memory. We precisely model the costs and constraints corresponding to Deep Learning frameworks such as PyTorch or Tensorflow, we propose optimal algorithms to find a valid sequence of memory-constrained operations and finally, we evaluate the performance of proposed algorithms on realistic networks and computation platforms. Our experiments show that the possibility to offload can remove one third of the overhead of rematerialization, and that together they can reduce the memory used for activations by a factor 4 to 6, with an overhead below 20%.

Efficient Combination of Rematerialization and Offloading for Training DNNs
Olivier Beaumont
Lionel Eyraud-Dubois
Alena Shilova
Inria Bordeaux
folivier. beakumnt,, lionel. mawnaleeyraudddubbo.e alema. shilova)@innia fr
Abstract
Rematerialization and offloading are well known strategies save memory during the training phase of deep neural networks, allowing data scientists to consider larger models, batch sizes or higher resolution data. Rematerialization trades memory for computation time, whereas Offloading trades memory for data movements. As these two resources are independent, is appealing to consider the simultaneous combination of both strategies to save even more memory. We precisely model costs and constraints corresponding Deep Learning frame- works such as PyTorch or Tensorflow, we propose optimal algorithms to find a valid sequence of memory-constrainn formance of proposed algorithms realistic networks and computation platforms. Our experiments show that the possibility to offload can remove one third the overhead of rematerialization, and that together they reduce the memory used for activations by a factor to 6, with an overhead below 20%.
operations and finally, we evaluate the per-
Introduction
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 GPT-3 30 have billions of parameters
34 or SENET [S] routinely have hundred million parameters, which raises difficult memory issues on architectures like GPUs. Moreover, it was observed [29] that sufficiently large batch-size is required to obtain good and fast convergence. In this paper focus 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 33 27 11], model parallelism (25] 37 26). filter or kernel [ or image parallelisms [8]. The combination with parallel strategies is however out of the scope of the present paper and goal here to model, analyze the complexity and understand how to benefit simultaneously from rematerialization and offloading.
and vision networks like EfficientNet (35], AmoebaNet
Rematerialization 14 20 3 is technique that consists in deleting from memory some of activations computed during the forward phase 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 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 the time the operation) consisting elementary forward, backward and delete operations, and which has minimum execution time among all sequences that satisfy given memory constraint. a technique that has the same objective, iee. to produce a sequence of operations satisfying 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
actually in memory at
Offloading (32 23
35th Conference on Neural Information Processing Systems (NcurIPS 2021)

---

the CPU and the GPU for lower memory requirements. The objective, in the context of offloading, 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 and model and notations in Section 3 The algorithmic solution to combine optimally checkpointing and offloading is presented in Section (nnd and experiments to assess the advantage of using both offloading and rematerialization, based on large range of neural networks and architecture parameters, are
presented in Section
2 Related work
Parallelism- memory optimization 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 [3 27 M 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 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 needed in order 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.
25 37 is used to distribute weights over different available
Remattrialization-bass strategies Rematerialization techniques consist in keeping only a subset of activations in memory during the whole training phase, while the others are dynamically recom- puted 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, com- putational 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 [33
The use of rematerialization strategies advocated for DNN in several papers 14 implementation of a simple periodic strategy [5] was provided in PyTorch [28] 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
5 20 An
The case of non sequential networks was also addressed in the literature 19 mostly under the simplifying assumption that discarding computed activations is only possible before the computation of the loss function 22 [0]. In [9] the authors proposed an algorithm
18
22

---

@
a
a L11
F1.-1
1088
@
a²
7
a
a
c++
B
B2
Bs
B₁
Bin
8
8
8
82
8
=
Figure 1: Data dependencies induced the training phase of Sequential Deep Neural Networks.
dealing with dynamic data flows, though based on an heuristical approach. The authors of 18] considered minimization of the runtime with the 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.
Oflladdig-basee strategies Offloading 32 is alternative approach which consists in of- floading 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 simple mechanism of memory virtualization, that introduces unnecessary synchronizations between data transfers and computations of later forward activations. This technique strategies to decide which activations to offload into the main memory are relatively straightforward, and basically consist in sending the CPU either all activations only those of convolutional layers. In [), the complexity of the offloading problem was analyzed and several NP-completeness results provided, depending whether activations be offloaded partially or not.
improved in In both papers, the algorithmic
In order to reduce the overhead incurred by the communications, the use of compression was advocated in [33). Memory virtualization
further considered in (24 23 16 38). Finally, the speed-centre rematerialization from [5] enhanced with memory-centtie rematerialization (discards combined with the simple offloading approach from [2 in [36] In addition, the authors [3 also used rematerialization of [5] with possibility of further offloading saved checkpoints to the CPU if rematerialization only is not enough to perform training under memory constraint.
activations of every segment all the time)
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 [ that corresponds to the processing of mini-batch This chain depicted by Figure i defined by the set of operations F; and By. The default forward operations Fg requires only the activation produced by the previous layer Fi-1 (or the input mini-batch in the case of Fi), whereas the default backward operation B; require the preceding gradient 5², the activation a E-1 and the intermediate generated data â' necessary for the backward step. The distinction between a and ā' is crucial in practice, since it allows consider computation graphs which have a sequential structure, but are not purely sequential. In this setting, it is indeed possible for F5 represent complex operation (any Direct Acyclic Graph of layers), and thus ã contains all the intermediate activations produced by F, and necessary to compute Br, whereas a only contains the output activation of Fz, that will be used by F241- In general, al can therefore be much smaller than a!
In order manipulate different types of activations, define in Table Ta a set of generic forward and backward operations, that can be directly translated in automatic diffcrenniation packages such as tf. Gradi aadientTape for Tensorflow or torch. autograd for PyTorch. Operations are defined both by the type of operands they use and produce (a' or â') and by the treatment they make of the input data after processing (deleted or kept in memory). Each operation is associated to processing time and a memory overhead that takes place during the computation. Both the processing time and the memory overhead only depend whether the operation is forward backward task.
In addition to the operations described in Table T we also consider two operations Or and P. that correspond respectively to offloading to and prefetching from CPU memory. OF and Pa can operate

---

Operation
Input
Output
Time
Memory overhead
{a
{a--1 a²
Fa
Forward and save all
{a
{at-1 @²)
UFF
{at-1j
{at-1 at}
Fd Forward and materialize input
fā--1}
far-1 ot}
UFF
of
fat-1)
{at}
Fg Forward without saving
{a--1j
{ut]
uFr
for, a. a 1
{0
B
Backward step
f8r, af a'-1)
fo
a--13
uB
o8
Table I: Generic operations available in DL frameworks.
both a-1 and @f-1, depending on which type of activation has been produced. Thanks CUDA streams, is possible to asynchronously offload and prefetch activations from and to the CPU memory, while independent computations are performed.
3.2 Assumptions and limitations
In what follows, our goal is to find a sequence of Flu-- Be, F F O2 and Pt operations to compute 8 from a". Given a bandwidth / between the GPU and the CPU, we can compute the peak memory usage and the total duration of a particular sequence. Our goal is find sequence (i) whose memory peak is smaller than the available memory on the GPU MaPu, and (2) whose overall duration is minimal. In order make the problem tractable, we introduce additional assumptions
1. Offloading before loss: all offloading operations must complete before the computation of the loss and prefetching operations (each prefetching operation being performed only once) start only after the computation of the loss. This is a reasonable assumption since offloading (resp. prefetching) introduce delays and therefore should be performed early (resp. as late) as possible.
2. Memory persistency holds: if for some layer i activation a' is computed, then operations related to layers i' with i' take place until B;. This is a common assumption in the rematerialization literature, and it was shown in [3] that in all practical cases, optimal solutions satisfy this assumption (despite theoretical result showing that there are cases where this is not true)
3. Asynchronous and continuous offloading: activations should be offloaded/prefetcheed en- tirely, but memory allocation and release can be performed in several steps during the communication. This assumption proved be efficient in provide a tractable context without degrading the solution quality, and indeed without this assumption, the problem becomes strongly NP-hard. Continuous offloading thus allows to write the dynamic program to obtain efficient solutions; however these solutions are evaluated in Section i context where the tensors are discarded entirely and once.
In Section our goal is to find the optimal sequence of operations under above assumptions. Removing some of these assumptions may therefore improve the solution (by seeking the solution in a larger space), but priori at the cost of more expensive computations.
1. A first limitation comes from the ability to offload activation, which is not easy to implement in frameworks like PyTorch, because it requires modifying and accessing internal structures. Solutions that currently implement offloading, such as VDNN B2] 2]. therefore require thorough rewriting of the transfer operations.
2. A second limitation of our work is that it relies on the possibility of linearizing networks, i.e. representing them as sequence of operations, in which each operation can in turn be arbitrarily complex (and associated to any type of dependences). In practice, most networks are easy to sequentialize. the case of rematerialization only, optimal solutions based integer linear programming such Checkmate [18] were proposed. However, the cost

---

of integer linear programming limits their use to very small networks (smaller than those considered in Section [5) and their relaxation come without any guarantee. The extension of the proposed approach to complex dependency structures is nevertheless not straightforward.
4 Combination of offloading and rematerialization
We propose a dynamic programming to optimally solve the problem of finding a sequence containing rematerialization and offloading operations of minimal duration. This algorithm is denoted by pofo, for "Persistent with Offloading during Forward Only". To save space, we will not detail in the main paper all the equations of the dynamic program, which involves a large number of cases. We will focus in the main part of the paper on the intuitions and the general working principle of the dynamic program and refer the reader to Appendix B for detailed derivations and proofs.
4.1 Rationale of the different operations
The pofo algorithm is based on a sequence of choices, that consist in deciding, for each S S L (i) which type of operation F; we are going to use and whether we are going to compute at or a' (ii) whether we are going keep the input value at-1 or at-1 in the memory of the GPU, offload it in the memory of the CPU or completely delete it from the memory and (iii) how to compute B₁.
Concerning (i), the size of a' is in general larger than a'. Both can be used to compute Fi+1 and Bi+1, but if a² only is kept in memory (either CPU or GPU), then F; will have to be recomputed during the backward phase before B, Concerning (i), if we delete a²-1 or a' from memory, then will be necessary to recompute F;-1 before processing B,, but it will save memory for the whole subsequent sequence F,+1- ... F2 By ... B,+1. af or a' is offloaded on CPU memory (which may take some time), the memory will be available for the subsequent sequence, and a² or 1 will have to be prefetched before computing B, The interest of offloading obviously depends on the bandwidth 33 of the PCI bus. Concerning (iii), if the input activation (a' or a a has been deleted from memory, we will have to recompute it, starting from the last kept activation (either at or at for some k < i 1). Several rematerialization sequences are possible to do this operation that have potentially different durations and that offer different prefetching possibilities, depending on their memory profile, and the choice of the optimal sequence will depend on the memory state before computing B,, as detailed in Appendix
4.2 Intuition of the overall scheme and state variables
The dynamic programming solution is built in several steps. We consider separately the forward propagation within which the offloading of activations can be done. The forward propagation is followed then with the loss calculation synchronized with the end of offloading operations. After loss computation, the backward propagation interleaved with the prefetching starts.
The general principle of pofo is the following. During the forward phase (before the computation of the loss), we consider the layers in an increasing order (from 1 to L). At each step, we have to decide which operation to implement among F2 Fe Fg and in the case where the input activation is kept, we have to decide if we add O,, Inspired by the ideas of Automatic Differentiation [13] (but proposed in the context of a fully homogeneous chain without offloading), our dynamic programming relies on the following remark: any solution can be decomposed into parts, where each part computes layers to j (j 2 i), and only the input to layer i is stored in memory. Starting from a layer i, whose input needs to be saved in memory, pofo needs decide whether the input will be saved with F or Fi It also needs to decide the index j of the end of the corresponding part in the solution, so that j + starts the next part of the sequence, and its input will be the first activation kept in memory after i. A recursive call to the sub-chain starting from j allows to obtain the corresponding running time. Between the two saved activations of layers and j. a sequence of F/ takes place. As forwards Ft, and F correspond to two different behaviors, they generate two different cases of the dynamic programming that are considered in Appendix B. In addition, made for the input of layer i, either store it in the memory of the GPU or offload it to the CPU; this decision impacts the memory available for the rest of the chain and idle times due communications.
offloading decision needs to be
In order to use dynamic programming, we need to define set of variables that describes the state of the system at any instant. This set of variables should be as small possible, since it has direct
5

---

influence on the size of the data structure and on the computing time to solve the dynamic program. On the other hand, these variables must be chosen wisely and they must contain enough information to make decisions for subsequent layers and be updated according to these decisions. To evaluate correctly memory constraints and to compute idle times, is important to know how memory may vary between layers i and] j with saved inputs: from some minimal memory occupation till maximal memory occupation. Both can be described for the forward and the backward phases using only three state variables: in addition of the index of the last layer with saved input discussed above, we use A,, Ar, and A. in addition to a boolean variable x.
A, denotes the total GPU memory occupied by the saved values among a.,..., a²-2 and a' a!.....a' a that are not transferred the CPU.
AF, denotes the amount of data from a², ,..., a'-² and a² .....
a'-2 a' that the schedule still
needs to offload after F;-1-
B, denotes the amount of data from a², aa, ..... a and al ....., a that the schedule should prefetch before starting B;-
x is boolean specifying whether the input of layer i is a'-1 (if 0) or a'-1 (if 1).
We denote as I (1 z)a'- + za'-1 the size of the input activation for layer i. With these state variables, we can compute M the memory on the GPU after executing Fi-1, and M B. the memory the GPU before executing B (excluding 81-1) are given by M.. = A, + A I and M = A, + AB. +
These memory values represent the maximal memory occupation between layers i and j for the forward and backward phases respectively, and are used for computing idle times when overlapping communications with computations, according to Lemma [ On the other hand, to estimate the feasibility of the scheduled operations, we also need to know the memory after everything has been offloaded, which is either A; + IF (when I stays the GPU) or A; (when Ií" is sent to the CPU), which we use to obtain the maximal memory available on the GPU.
The updates of the state variables follow simple rules described below
if no new data is offloaded (prefetched), then Ar, (A),) is constantly decreasing from index to index j at speed B, until reaching zero;
if new data has be offloaded (prefetched), then the data size is added into AF, (A );
A is updated if new data is saved on the GPU (without being later offloaded to the CPU).
Our goal is to find the sequence of operations that minimizes the overall execution time. However, both rematerialization and offloading can induce extra time with respect to the execution of the sequence F1... ... FrBr... ... B₁ which can be computed given infinite memory on the GPU:
During the forward phase, offloading helps keep GPU memory low, but, if the transfers to the CPU are not fast enough due to limited bandwidth, some idle time might occur, waiting for enough memory be freed by offloading. This is analyzed in details in Appendix [B.
At the interface between the forward and the backward phase, we need to enforce that offloading and prefetching are well synchronized by the computation of the loss, what can in turn introduce some idle time on the GPU. This will be analyzed in Appendix B.2
During the backward phase (after the computation of the loss), there might be two sources of delays: recomputations of discarded activations and idle times induced by prefetching. Indeed, some prefetching operations may not be completed by the time the activation is needed, if they are delayed because of memory constraints. In fact, these two sources of extra times are not independent, since a longer rematerialization sequence with lower memory needs might allow more overlapping of prefetching with computations and thus avoiding idle time later. We propose an auxiliary dynamic program that includes the computation of intermediate idle times in each recursive call to find the best schedule under prefetching. The idle times for prefetching are found using Lemmal) I This is analyzed in Appendix B.3

---

Finally, we present a lemma that describes the general scheme for optimally overlapping the commu- nications with scheduled operations. Lemma iss valid for both prefetching and offloading, as their behavior is symmetrical. Offloading takes place at the beginning of the sequence, making available memory increase at the speed from its initial value mmmin- On the contrary, prefetching is done at the end of the sequence, making available memory decrease at the speed 3 until reaching miin Performing offloading as soon as possible and prefetching as late as possible allows to have more available memory for the middle of the execution. Note that mimi could be equivalently replaced with MGPU - Mmax, where Mmax denotes the maximal memory occupation (either M;, or M,))
Lemma Let us consider a fixed sequence of operations. for which the available memory increases (resp. decreases) during execution because of data unloading (resp. prefetching), with its minimum at momin- Let us denote by M the memory required to process operation E S, and d, the distance between o and the end (beginning) of sequence S,i.e the cumulative duration of operations taking place before (after) operation o. Then, the execution of S needs to be delayed by some idle time
(maxoes(M) - -Bdo) - -mmsn 0)
= max
8
The above mentioned ingredients can be used to build the dynamic programming algorithm pofo, solving optimally the problem in the pseudo-polynomial time. Theorem I states this result, while the complete description of the dynamic program and the proof of the theorem can be found in Appendix B of the memory values. In the experiments of Section 5 we systematically use 50 discretization steps (scaling so that MGPU = 50), which shows no discretization effect (we tried more steps to confirm this) and a reasonable complexity, with an execution time below 4 minutes in the worst cases. Since optimization performed only once before the training starts, we argue that this is fully acceptable.
The algorithm
in the memory limit and it relies on discretization
Theorem Under assumptions of Section 3.2 the problem of finding the minimal processing time for using operations from rogether with offloading O₂ and prefetching Pe. under memory limit McPU discretized with NoPU values and bandwidth B, can be solved optimally with dynamic programming algorithm with a complexity of O(L (2 + L3 N2 N&Pu).
the chain from Figure [
5
Experiments
5.1
Additional heuristics
We have implemented two other sophisticated heuristies for comparison purposes. These heuristics not bound with the assumptions of Section 3.2, but come without any optimality for the produced sequences.
In the first heuristic, called opportunist, the objective is to use the communication medium as much as possible. We compute the first layer with Fan and offload its input and output, ensuring that the memory will remain fully available for the rest of the computation. The next layers are computed with Fg until the end of communications. We then start new communication by performing the next layer with Fal and offloading its input and output, and so on until the end of the sequence can be entirely performed memory. This process thus builds groups of layers between two Fat operations. We then compute the backward phase for each group using the rematerial iemattriaaizatiinn algorithm, and concatenate them (with the necessary prefetches) to obtain the final sequence. Note that opportunist is conceptually close to the implementation from Superneurons [36..
The second heuristic is called autocapper. It relies on an internal capper algorithm, that uses an in- creased memory limit M' > Maru as an additional input. capper computes a pure rematerialization sequence with the limit M', finds the peak memory usage in the sequence and offloads the lowest indexed activation present in GPU memory at that time. This process is repeated until the sequence fits in the memory limit MGPU. autocapper calls capper with 40 values of M' evenly spaced between the target MGPU and Mbigh, the memory required without rematerialization or offloading. For each value M', the resulting sequence from capper is simulated and the best one is kept. Note that, when autocapper does not perform recomputations, it behaves as the GREEDY heuristic in 4
7

---

densenet 121
densenet 161
densenet 169
inception 200
Overhead wrt sequential (for 12GB's bandwieth)
1.3
1.2
Algorithm
aulocapper
offiead
resnet 101
resnel 152
resinet 200
resinet 50
opportunis:
1.3
POFO
1.2
remaaerializzsion
1.0
022
0.4
0.6
0.2
0.4
0.6
0.2
0.4
0.8
0.2
0.4
Memory Ratio
Figure 2: Simulation results for fixed bandwidth 8 = 12GB/ss and varying memory ratio a.
5.2 Simulation results
We measured running times and memory occupation of several networks from PyTorch torchvision package: resnet, densenet and inception, with batch size of 16 and images of 500 500 pixels. For simplicity, we only present single value for batch size: we also experimented with other batch and image sizes, and obtained very similar results. Time measurements were performed a NVVidia Tesla V100 GPU. We also measured the bandwidth obtained when transferring Py Torchh tensors from and the GPU and obtained 12GB/s. The simulation results presented here were obtained using 4 cores of a 24-core Haswell Intel Xeon E5-2680 v3 at 2,5 GHz, with 128GB of memory, and used about one hour of computation. We modified the open-source rematerialization framework rotor to handle offloading, and we implemented our algorithms in this framework.
We consider 5 algorithms in total: our dynamic program pofo, the optimal rematerial inttterriiization. only algorithm from [3, the optimal offload-onl approach (DyNPRoG heuristic from (4), and autocapper and opportumist heuristies. We the sequential time reference: is the time that it would take to process forward and backward phases with infinite memory, i.e. the sum of all forward and backward times. For each network, compute the highest and lowest memory requirements (denoted respectively high and Allw):: high is obtained with the sequential i approach, while M obtained by recomputing everything from the beginning each step of the backward phase. We can thus explore the whole range of achievable memory sizes for this network, by considering values within the interval [M .M for a given ratio [0, 1]. the memory limit Maru is set to (1 a)Ml + high We consider values of a from 0.05 (low memory scenarios) to 0.8 (high memory scenarios). Larger values of are not included since this makes the memory limit so large that the optimization problem becomes easy, and all algorithms behave very similarly, with almost zero overhead. Depending on the networks, the values of Mlow range from to 2.5GB, and the values of hihggh range from 4 to 20GB. On the other hand, the sizes of the weights range from 50 250MB, which highlights that the main memory usage for these networks comes from the activations. For the algorithms which perform offloading, we also vary the bandwidth value B, from 12GB/s (a realistic scenario) to 36GB/s (corresponding to possible future improvements).
The results are shown on two plots: Figure 2
presents the behavior of the algorithms for fixed bandwidth and varying memory constraint, while Figure 3 explores the effects of increasing the bandwidth for a given memory limit. We draw the following conclusions:
When memory is large (a = 0.75), offloading more effective than rematerialization: with 3 = 12GB/s, there is time to offload and prefetch enough data to run the whole chain. However, for smaller memory limits, the offl oadlooyyyyy policy exhibits much worse performance than rematerialization (unless more bandwidth is available, see Figure 3).

---

densenet 121
densenet 161
densenet 169
inception 200
Overhead wrt sequential (for memory railo 0=0.
1.4
1.2
Algorithm
1.1
aulocapper
officad
resnel 101
resnel 152
resinet 200
resine 50
1.4
POFO
rematerializzsion
1.2
1.1
1.0
18
24
36
12
36 12 16
24
24
Bandwidth (GB/s)
Figure 3: Simulation results for fixed memory ratio a = 0.2 and varying bandwidth.
Unless in cases with small memory, the opportumist policy achieves slightly better performance than pure remater ierriall izaatiin,,, but is sometimes unable to produce a solution when the memory limit is too low. On the other hand, the more sophisticated autocapper algorithm obtains very good performance for memory ratios above 0.5.
The pofo algorithm successfully combines the advantages both rematerialization and offloading, and consistently outperforms both of them. When the memory limit is high enough, the optimization problem is relatively easy, and autocapper and pofo achieve similar performance. In cases autocapper is marginally better than pofo, pofo produces solutions under assumptions of Section 3.2 while partial memory releases (as- sumption (iii) are difficult implement in practice. For lower memory limits, the more optimized pofo algorithm is able to produce much better sequences.
As can be seen on Figure 3 the offload-only approach works perfectly when the bandwidth is high. Indeed, if there is enough communication capability to offload all the data while it is produced, it is possible to avoid recomputations without inducing idle time. Compared to pure offloading techniques, our pofo solution allows to train models with large activations with cheaper communication links.
more precise analysis shows that the solutions computed by pofo offload significant amount of data, about 25% of the total activation size for the inception and resnet networks, and 20% for densenet. This represents for example GGB of offloaded data for reanet 101. For the same example, the amount of discarded data (which is recomputed later through rematerialization) varies from (for a 0.75) to 1.3GB (for a = 0.1). general, the reason why pofo more efficient than offloading alone is that rematerialization reduces memory usage, which means that fewer activations need be offloaded. This allows fit the offloading and prefetching time within the computing time. Symmetrically, the reason why pofo is more efficient than rematerialization alone because offloading allows to remove some activations from the GPU memory "for free" in terms of recomputation. The rematerialization procedure thus needs to recompute fewer layers, which reduces the overhead cost.
5.3 Experimental results
In the recent release of Py Torch 1.10, the introduction of the serdd__ensors_hoooo () feature makes it possible to implement the offloading technique described in this paper. Indeed, this function allows to register hooks that can capture all tensors generated by an operation and then "pack" them (compress or offload) during the forward propagation and "unpack" them (extract or prefetch) during the backward propagation. We have implemented preliminary version of our best performing

---

densenet 121 BS=32
densenet 161 BS=32
densenet 169 BS=32
inception 200 BS =64
1.25
1.20
1.15
Overhead wrt bestbetttee solution
1.10
Algorithm
1.05
1.00
autocapper
6
10
12
6
10
12
6
10
12
14
7.5
10.0
12.5
offload
resnet101 B5=32
resnet B5=16
resnet 200 B5=16
resnet 8S=32
POFO
1.25
rematerielization
1.20
sequendial
1.15
1.10
1.05
1.00
10
12
2.5
5.0
7.5 10.012.5
5.0 7.5
10.0
12.54
6
8
10
12
Memory used (GB)
Figure 4: Experimental results on actual runs.
algorithms (pofo and autocapper) and made them available rotor M. We further present the current results in Figure obtained by testing these new extensions of rotor on the same neural networks as in the previous section, this time on an actual V100 GPU.
Overall, these preliminary experiments confirm that combining offloading with rematerialization allows to significantly improve over pure rematerialization in most scenarios. In particular, pofo is still the best algorithm for memory optimization among the ones considered in the plot. In most cases, pofo either shows the smallest overhead or behaves in the same way as other methods. However, densenet- 169 is the exception to this trend, where rematerial wwriiaaalittttooioil is strictly better than pofo or autocapper. but even in this case pofo outperforms autocapper. On the other hand, the case of resnet-100 101 demonstrates the significant improvement of pofo over other strategies, highlighting that there exist cases for which the benefit is important.
This discrepancy in the results obtained in simulation and in the experiments comes from the fact that the current implementation of pofo and autocapper in rotor uses more memory than expected. Our simulation results hint that further optimization of this preliminary implementation should allow obtain even better performance.
Conclusion
In this paper, we formalize the problem of the optimal combination of rematerialization and offloading. which are two classical strategies for coping with memory limitations on a GPU. We show that the optimal solution can be computed using dynamic programming. in few seconds minutes for very deep networks. From experiments, we show that the combination of offloading and rematerialization is very efficient and allows, many cases, to transparently perform training with to times less memory, at the cost of a 10-20% time overhead. Among the perspectives of this work, one could be interested minimizing the energy consumption under both memory and throughput constraints. This study would enable to quantify the relative effects of data transfers and redundant computations from a different viewpoint. Another perspective of this work would be to adapt the dynamic programming approach to the case where we consider the possibility of offloading not only activations, but also the weights of the network, which is necessary in the case recent very large NLP networks in particular and which is used, for example, on the basis of heuristic optimization in Deepspeed 3] The dynamic programming approach also seems to adapted to this new context. Finally, it could be possible extend the results of this paper to more general memory hicrarchies, which include the case of CPU with limited memory and a possibility of offloading to disk.
few
http::/wwwwweepppeeee
10

---

Acknowledgments and Disclosure of Funding
We would like to thank our reviewers for their useful feedback and insights helping to improve our work. We would like to express special thanks for pointing out the new useful function of Py PyTorch that made our implementation possible. This work was supported by The HPC-BiiData Inria Project Lab (IPL).
References
[1] Rotor: httppp//ghhbb.....rrrrr/aaaaaaop
2019.
[2] Shriram S B, Anshuj Garg, and Purushottam Kulkarni. Dynamic memory management for gpu-based training of deep neural networks. In IEEE International Parallel and Distributed Processing Symposium (IPDPS). IEEE Press, 2019.
[3] Olivier Beaumont, Lionel Eyraud- Eyraud-Dubois. Julien Herrmann, Alexis Joly, and Alena Shilova. Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory. Research Report RR-93 RR-9302, Inria Bordeaux Sud-Quest, November 2019.
[4] Olivier Beaumont, Lionel Eyraud-Dubois, and Alena Shilova. Optimal gpu-cpu offloading strategies for deep neural network training. Maciej Malawski and Krzysztof Rzadca, editors, Euro-Par 2020: Parallel Processing, pages 151-166, Cham, 2020. Springer International Publishing.
[5] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174, 2016.
[6] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Andrew Senior, Paul Tucker, Ke Yang, Quoc Le, et al. Large scale distributed deep networks. In Advances in neural information processing systems, pages 1223-1231, 2012.
[7] Jacob Devlin, Ming- Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: (0t::/8/0.0000 00.04805, 2018.
[8] Nikoli Dryden, Naoya Maruyama, Tom Benson, Tim Moon, Marc Snir, and Brian Van Essen. Improving strong-scaling cnn training by exploiting finer-grained parallelism. In IEEE International Parallel and Distributed Processing Symposium. IEEE Press, 2019.
[9] Nikoli Dryden, Naoya Maruyama, Tim Moon, Tom Benson, Marc Snir, and Brian Van Essen. Channel and filter parallelism for large-scale enn training. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, page 10. ACM, 2019.
[10] Jianwei Feng and Dong Huang. Optimal gradient checkpoint search for arbitrary computation graphs, 2018.
[11] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. 06 2017.
[12] Andreas Griewank. On automatic differentiation. Mathematical Programming: Recent Devel- opments and Applications, 6(6):83-107, 1989.
[13] 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.
[14] Audrunas Gruslys, Rémi Munos, Ivo Danihelka, Marc Lanctot, and Alex Graves. Memory- efficient backpropagation through time. Advances in Neural Information Processing Systems, pages 4135-1133, 2016.

---

[15] Jie Hu, Li Shen, and Gang Sun. Squeeze-ande and-excitation networks. In Proceedings the conference on computer vision and pattern recognition, pages 7132-7141, 2018.
the IEEE
[16] Chien-Chin Huang, Gu Jin, and Jinyang Li. Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping. In Proceedings of the Twenty-Fifth International Conference on Architectural Support for Progrumming Languages and Operating Systems, pages 1341-1355, 2020.
[17] Yanping Huang. Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoe V Le. Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Advances in Neural Information Processing Systems, pages 103-112, 2019.
[18] Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Kurt Keutzer, Ion Stoica, and Joseph E. Gonzalez. Checkmate: Breaking the memory wall with optimal tensor rematerialization, 2019.
[19] Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, and Zachary Tatlock. Dynamic tensor rematerialization. arXiv preprint 1oXi::2006.09666 2020.
[20] Navjot Kukreja, Jan Hückelheim, and Gerard J Gorman. Backpropagation for long sequences: beyond memory constraints with constant overheads. arXiv preprint arXiv:/80k 0f 17, 2018.
[21] Ravi Kumar, Manish Purohit, Zoya Svitkina, Erik Vee, and Joshua Wang. Efficient rematerializa- tion for deep networks. Advances in Neural Information Processing Systems, 32:15172-151811, 2019.
[22] Mitsuru Kusumoto, Takuya Inoue, Gentaro Watanabe, Takuya Akiba, and Masanori Koyama. A graph theoretic framework of recomputation algorithms for memory-effieiene backpropagation. arXiv preprint arXiv: 19X::1905.11722. 2019.
[23] Tung D Le, Haruki Imai, Yasushi Negishi, and Kiyokuni Kawachiya. Tflms: Large model support in tensorflow by graph rewriting. arXiv preprint 1007:v8807000000 2018.
[24] Chen Meng, Minmin Sun, Jun Yang. Minghui Qiu, and Yang Gu. Training deeper models by gpu memory optimization on tensorflow. In Proc. of ML Systems Workshop in NIPS, 2017.
[25] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. PipeDream: generalized pipeline parallelism for DNN training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles, pages 1-15, 2019.
[26] Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory- efficient pipeline-paralle dan training. arXiv preprint arXiv: 09500,, 2020.
[27] Thomas Paine, Hailin Jin, Jianchao Yang, Zhe Lin, and Thomas Huang. Gpu asynchronous stochastic gradient descent to speed up neural network training. arXiv preprint arXiv: 2013.
[28] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary De De Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch, 2017.
[29] Bharadwaj Pudipeddi, Maral Mesmakhosroshahi, Jinwen Xi, and Sujeeth Bharadwaj. Training large neural networks with constant memory using new execution algorithm. arXiv preprint a0Xiv:2002.05645, 2020.
[30] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 18)): 2019.
[31] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimiza- tions toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-16. IEEE, 2020.
12

---

[32] Minsoo Rhu, Natalia Gimelshein, Jason Clemons, Arslan Zulfiqar, and Stephen W Keckler. vdnn: Virtualized deep neural networks for scalable, memory-efficien neural network design. In The 49th Annual IEEE/ACM International Symposium on Microarchitecture, page 18. IEEE Press, 2016.
[33] Minsoo Rhu, Mike O'Connor, Niladrish Chatterjee, Jeff Pool, Youngeun Kwon, and Stephen W Keckler. Compressing dma engine: Leveraging activation sparsity for training deep neural networks. In 2018 IEEE International Symposium on High Performance Computer Architecture (HPCA), pages 78-91. IEEE, 2018.
[34] Syed Asif Raza Shah, Wenji Wu, Qiming Lu, Liang Zhang, Sajith Sasidharan, Phil DeMar, Chin Guok, John Macauley, Eric Pouyoul, Jin Kim, et al. Amoebanet: An sdn-enabled network service for big data science. Journal of Network and Computer Applications, 119:70-82, 2018.
[35] Mingxing Tan and Quoc Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In International Conference on Machine Learning, pages 1105-6114 PMLR, 2019.
[36] Linnan Wang, Jinmian Ye, Yiyang Zhao, Wei Wu, Ang Li, Shuaiwen Leon Song, Zenglin Xu, and Tim Kraska. Superneurons: Dynamic gpu memory management for training deep neural networks. SIGPLAN Not., 53(1):41-53, February 2018.
[37] Jun Zhan and Jinghui Zhang. Pipe-torch: Pipeline-based distributed deep learning in gpu cluster with heterogeneous networking. In 20/9 Seventh International Conference on Advanced Cloud and Big Data (CBD), pages 55-60. IEEE, 2019.
[38] Junzhe Zhang. Sai Ho Yeung, Yao Shu, Bingsheng He, and Wei Wang. Efficient memory management for gpu-based deep learning systems. arXiv preprint arXiv:1903.06631, 2019.
[39] Martin Zinkevich, Markus Weimer, Lihong Li, and Alex Smola. Parallelized stochastic gradient descent. In Advances in neural information processing systems, pages 2595-2603, 2010.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section 3.2
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] in the Appendix.
3. If you experiments..
(a) Did you include the code. data, and instructions needed reproduce the main perimental results (either in the supplemental material or as URL)? [Yes] in the supplemental material
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [N/A]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [N/A]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes] in Section 5
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] rotor, cited in Section
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either the supplemental material or as URL? [Yes] in supplemental material
(d) Did you discuss whether and how consent obtained from people whose data you're
using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "spliLGq1lP", "reviewer": "Reviewer_39eR", "summary": "The paper presents an algorithm, called pofo, for optimizing the execution time of DNN training on a single GPU, considering a combination of rematerialization and offloading strategies.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Major Comments\n1.\tRematerialization and offloading have been already proposed in the literature. It is not clear which are the challenges of combining both approaches, if any. Therefore, the contributions do not seem much novel.\n2.\tIn practice, can the proposed framework adapt to a limited memory size as a given constraint?\n3.\tDoes the proposed method work only with feed-forward networks, or also with recurrent networks?\n4.\tIn Section 5.2, “In some cases autocapper is marginally better than pofo, as pofo produces solutions under assumptions of Section 3.2, while partial memory releases (assumption (iii)) are difficult to implement in practice.” Why are there different assumptions for the autocapper procedure? For comparison purposes, the assumptions should be the same for every algorithm. Moreover, the reasons why the autocapper outperforms the proposed pofo should be discussed more comprehensively.\n\nMinor Comments\n1.\tThere is no need to report Figure 1 and Table 1 in the appendix, since they are already reported in the main manuscript.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "PgAkB5ieuYk", "reviewer": "Reviewer_ZiAJ", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper tackles a very important problem because GPU memory is perhaps the biggest bottleneck in DL since it grows at significantly slower rates compared to model, data, and compute.  Also, addressing activation memory, which is one of the major sources of memory consumption, is very impactful to the community. A solution that consists of optimally combining SOTA techniques is very appealing and the automated nature could realistically save the massive wastage of human and machine resources on tuning SOTA techniques. \n\nThe paper is very well written and I found it enjoyable to read end-to-end. I felt the authors did such a great job in how they used available space fit enough high-level intuition and low-level details such that I could fully grasp the paper without referring to the appendix.\n\nConcerns:\n1. A key weakness is that evaluation does not consider multi-gpu scenarios, which are more common for large-model training than single gpu. This seems like an unnecessary omission that is not adequately explained. Moreover, since your results are simulation-based it does not seem like much extra work is required. It seems that all that is needed is to consider PCI bandwidths per GPU lower than 12GB/sec (e.g., 3GB/sec per GPU for 16 GPUs in a DGX-2 box). This way you won't need to adjust your algorithm to consider inter-gpu communication (which is typically fast within a node anyways). Please clarify if I am missing anything. \n2. I can't figure out if your results apply to language models. The key difference of language models that could affect the effectiveness of your approach is that model and optimizer consume much more memory leaving very little room for activations. On the other hand, and perhaps as a consequence, the batch size is typically much smaller than image models. So perhaps these two factors will cancel out and your results will still hold, but it is hard to tell since you don't have language model results. Another view is that perhaps your work is more desperately needed in the language model space where batch sizes are heavily constrained by memory pressure. Can you share any thoughts on this? Do the claims in the paper need adjusting to reflect the image-centric nature of the results?\n\nSuggestions/questions:\n1. The evaluation should include model and activation sizes, as well as the GPU memory size for the presented results.\n2. Offloading to NVMe is already possible using ZeRO-Infinity and Nvidia GPUDirect (either from CPU memory or directly from GPU memory). This would be needed when CPU memory is not large enough for the activations. Can your approach can be extended for this?      ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"ethical_concerns": "No.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "90n6CSp0-3F", "reviewer": "Reviewer_nNiv", "summary": "This paper tackles the problem of combining rematerialization with offloading by using dynamic programming, and demonstrate (in simulations) that their approach can substantially outperform pure offloading as well as pure rematerialization in terms of runtime in different situations.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I like the problem that this paper tackles - how do you balance various strategies for making your model bigger/faster? All of these different approaches have different tradeoffs, and it's often unclear how you should use one vs. the other. So, papers that attempt to resolve these questions are quite valuable.\n\nThe DP itself is fairly straightforward but reasonable.\n\nMy primary qualm with this paper is that the experiments are all performed with a simulator. In my opinion, the DP in this paper is not novel enough to have intrinsic novelty. However, the DP algorithm can still be a valuable contribution if they demonstrate that the DP actually works in practice. Much of the contribution of DP algorithms in situations like this is demonstrating that the assumptions made by the algorithm are realistic enough for the DP to result in a good solution. However, if you're simply running it on a simulator, then I'm not sure this successfully tests whether your algorithm's assumptions are realistic.\n\nRegardless, I think that the simulator results are quite interesting (with good baselines), and demonstrate that offloading may be a technique that people trying to reduce activation memory should consider looking into. \n\nTo be honest, I am quite split on this paper. On the one hand, I think that this paper tackles a valuable problem, demonstrates some novelty in combining rematerialization + offloading, and has intriguing simulator results. On the other hand, without experiments on real hardware or actual software, I am skeptical that these paper's results will have much impact in the community.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "7lCST_oGqv", "reviewer": "Reviewer_Zjqt", "summary": "This paper investigates the advantages of _combining_ the memory optimizations of rematerialization and offloading for training large models on single GPUs.  In doing so it studies when such a combination is useful and unsurprisingly shows that a combination of the two provides better benefits than either optimization used in isolation while incurring a modest 20% runtime overhead.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I enjoyed reading your paper and your focus on memory saving optimizations for activations to push the boundaries of single GPU training.  I was however quite confused by all the discussion on offloading weights (lines 109-115, references to deepspeed and such).  My suggestion: just explicitly state up-front that your focus is on activations (and you assume weights and optimizer state fit in memory) and just ignore all this distracting discussion on offloading of weights.\n\nI was a little surprised that when it came time to evaluate you picked a batch size of 16.  Could you provide a rationale as to why you pick this batch size for all the DNNs you evaluate with?  Is this something that you can fit for _sequential_ for each of the 8 DNNs you evaluate with? I assume you do this so that you can measure overhead with respect to sequential.  And if so, when you throw your memory saving optimizations at it, are you leaving memory on the GPU unused?  It is ok if you are doing this just for evaluation, but it would be good to call this out.\n\nIs there a reason why your graph x-axis (alpha) ends at 0.8?  Why not show us what happens at 0.9 and 0.95?\n\nOne thing that seemed like a missed opportunity in studying this combination of optimizations was reporting what percentage of activations were offloaded versus what percentage were recomputed.  I suspect that because offload is limited by PCIe throughput, offload is used sparingly and spread out by large temporal intervals so that prefetching these tensors can be hidden.  It would be great if you can provide some quantitative data and intuition for this.\n\nAs I am writing this review, I realized that I am assuming you are prefetching tensors that you had offloaded (and that you are not doing this fetch in the critical path).  When you offload, are you assuming that you have space buffers available to prefetch data into?  That is, from the total memory available to you are reserving some memory for prefetch buffers?  And is offloading an option only because you are using rematerialization in between - that is you are increasing compute time and hence providing more of an opportunity to hide prefetching from slower CPU memory over PCIe (12 GB/s)?\n\nOverall, while you show us that the combination of the two memory optimizations you consider are effective (compared to individual optimizations), it is unclear _why_ offloading is helping in this combination and how much offloading is being performed.  You have an opportunity to nail this analysis (even if it means suggesting offloading is used sparingly but should be used nonetheless in this combination).\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are mostly covered in my review above, but succinctly: the paper could do a better job explaining how offloading is used in conjunction with rematerialization - that is not just the effect of using them, but what is the sequence of offload/rematerialization operations are the algorithms in the paper suggesting (some sort of temporal view for intuition or percentage of tensors offloaded versus rematerialized for quantitative data).", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes an optimized algorithm to compute a sequence of forward / backward / offload / prefetch operations on activations that optimizes training throughput of linearized DNNs under memory constraints. The rebuttal solves the reviewers' concerns, and the reviewers unanimously agrees to accept the paper. 
Decision: Accept (Poster)