Dynamic Tensor Rematerialization

√ N ) memory budget with only O ( N ) tensor operations . DTR closely matches the performance of optimal static checkpointing in simulated experiments . We incorporate a DTR prototype into PyTorch merely by interposing on tensor allocations and operator calls and collecting lightweight metadata on tensors . 1 INTRODUCTION . As state-of-the-art deep learning ( DL ) models continue to grow , training them within the constraints of on-device memory becomes increasingly challenging . The memory demands of emerging models prevent their training on memory-limited devices ( such as specialized accelerators , low-powered embedded devices , or older GPUs ) and limit researchers ’ ability to explore memory-intensive architectures and training techniques . Checkpointing is one technique that enables training with models and batches that exceed on-device memory without modifying the model ’ s design . It is achieved by freeing some activations from memory and recomputing them on demand . Adapted from techniques in automatic differentiation ( Baydin et al. , 2015 ; Griewank & Walther , 2000 ; Siskind & Pearlmutter , 2018 ) , checkpointing in the DL context exploits the fact that intermediate activations for backpropagation dominate memory usage during training ( Sohoni et al. , 2019 ) but can be easily recomputed by replaying parts of the forward pass . Current DL checkpointing techniques ( Chen et al. , 2016 ; Jain et al. , 2020 ; Kumar et al. , 2019 ; Gruslys et al. , 2016 ) statically plan which activations to recompute offline , requiring an initial stage of model analysis . In this paper , we demonstrate that static planning is unnecessary for DL checkpointing . We present Dynamic Tensor Rematerialization ( DTR ) , a greedy online algorithm for heuristically checkpointing arbitrary DL models . DTR operates like a tensor-level cache : it collects metadata on tensors and operators as a model is trained and uses it to guide heuristics that choose which activations to free and later recompute . As a runtime system , DTR can utilize dynamically gathered information ( e.g. , measured operator costs ) . Additionally , its simple , cache-like approach requires no advance knowledge of the model or application , letting it immediately support arbitrarily dynamic models and applications featuring higher-order differentiation . For example , given a model with data-dependent control flow like TreeLSTM ( Tai et al. , 2015 ) , DTR ’ s runtime can simply evict tensors when memory runs out and rematerialize them as needed . By contrast , static planning techniques assume a static dataflow graph , which requires “ unrolling ” dynamic models and performing ( potentially expensive ) planning for every distinct input . ∗Equal contribution . †Paul G. Allen School of Computer Science & Engineering , University of Washington , Seattle , WA ‡OctoML , Seattle , WA §School of Computer Science , Carnegie Mellon University , Pittsburgh , PA This paper describes DTR ’ s design ( Sec . 2 ) and makes the following contributions : • We prove that DTR can train an N -layer linear feedforward network on an Ω ( √ N ) memory budget with only O ( N ) tensor operations ( Sec . 3 ) , which is within a constant factor of optimal and matches the offline bound of the Chen et al . ( 2016 ) static checkpointing technique . • We formalize DL model checkpointing as an online rematerialization problem and define a greedy algorithm parameterized by caching-inspired heuristics . In simulated trials our heuristic attains near-optimal performance on a variety of DL models ( Sec . 4 ) . • We implement a DTR prototype by making only modest modifications to the PyTorch framework , enabling training under restricted memory budgets for both static and dynamic models and demonstrating the ease with which our algorithm can be incorporated into an existing DL framework ( Sec . 5 ) . Note that techniques other than checkpointing , such as swapping tensors between devices , can also enable training under limited memory . In Sec . 6 , we discuss these approaches and how they could operate with DTR . 2 DYNAMIC TENSOR REMATERIALIZATION . We introduce Dynamic Tensor Rematerialization ( DTR ) , a thin runtime layer that intercepts tensor allocations , accesses , and deallocations and eliminates the need for ahead-of-time model analysis to support checkpointing . Figure 1 shows DTR ’ s high-level approach . When a tensor allocation occurs ( AllocateBuffer ) , DTR first checks if sufficient memory is available . If so , it generates a fresh tensor identifier , initializes its metadata for future recomputation , allocates the requested memory , and returns a new tensor . If not , DTR heuristically selects and evicts resident tensors until the requested allocation can be accommodated . Constant tensors ( loaded from external data ) can not be evicted since no corresponding operation rematerializes them . Upon tensor access , DTR first checks if the tensor is resident in memory . If so , it updates tensor metadata before returning the requested tensor . If the tensor has been evicted , DTR rematerializes it by replaying the parent operation that originally produced the tensor . Crucially , rematerialization can be recursive : if the arguments to an evicted tensor ’ s parent operation have also been evicted , then they must first be rematerialized . Rematerialization may trigger more evictions if memory is exhausted during the potentially recursive process . Upon tensor deallocation ( other than by evictions ) , the runtime is invoked again ( Deallocate ) , letting it update tensor metadata and eagerly perform profitable evictions . Assumptions . This description of DTR assumes that : tensors are accessed only by opaque operators ; tensors are either constants or produced by operators ; operators produce individual tensors ; and operators are pure ( deterministic functions of their arguments ) . Under this model , a training epoch is simply a sequence of tensor operations without any inherent requirement to recognize training-specific structure , like the transition to the backward pass . DTR will evict as many tensors as necessary to avoid running out of memory . If all inputs and outputs of a single operation can not fit into available memory , rematerialization will fail ; therefore , on a given model and input , there may be a threshold for the lowest budget DTR can support . The choice of heuristic can affect the likelihood of failure since different eviction choices can result in deeply nested rematerializations that require many tensors to remain in memory . Heuristics . DTR is parameterized by heuristics that guide its eviction choices . As in caching , DTR ’ s eviction heuristic dynamically predicts which resident tensors are least valuable . The choice of heuristic determines what metadata ( additional runtime facts ) must be tracked for each tensor and operator and thus affects DTR ’ s runtime overhead . In our evaluation , we consider a runtime system that tracks the following metadata for each tensor t : staleness , s ( t ) , the time since last access ; memory , m ( t ) , the size of the tensor ; and cost , c0 ( t ) , the time required to compute t from its parent tensor ( s ) . We observe that DTR ’ s metadata overhead is low relative to the cost of typical DL tensor operations . We propose a rematerialization-specific heuristic that balances staleness , memory , and cost , evicting the tensor t that is stalest ( least likely to be needed soon ) , largest ( saves the most space ) , and cheapest ( requires the least additional rematerialization if t is needed again ) . To capture the total amount of rematerialization required if t is evicted , we sum the costs over the tensor ’ s evicted neighborhood e∗ ( t ) , i.e. , the set of evicted tensors that would either need to be rematerialized to recompute t or would need t to be resident to be recomputed . We define the projected cost , c ( t ) , of rematerializing tensor t as c0 ( t ) + ∑ t′∈e∗ ( t ) c0 ( t ′ ) . Using this definition , we define our heuristic , which evicts the tensor minimizing hDTR ( t ) = c ( t ) / [ m ( t ) · s ( t ) ] . By including both forward and backward dependencies of t in e∗ ( t ) , hDTR penalizes creating long chains of evicted tensors ( and hence potential recursive rematerializations ) that could arise from t ’ s eviction . To illustrate evicted neighborhoods , suppose DTR is checkpointing the network shown in Figure 1 , where the resident tensors are { t0 , t2 , t3 , t6 } . Before node t7 is computed , we have e∗ ( t2 ) = { t1 , t4 } and e∗ ( t3 ) = { t1 , t4 , t5 } . Since each new eviction can expand a given tensor ’ s evicted neighborhood and each rematerialization can shrink it , dynamically tracking evicted neighborhoods can introduce further costs at run time . To decrease runtime overhead , we developed an approximation of e∗ using an undirected relaxation tracked by a union-find data structure that uses a constant-time approximation for splitting . We use this approximation to define heqDTR analogously ( Sec . 4.1 and Appendix C.3 contain details ) , which performs nearly as well as hDTR in our evaluation ( Sec . 4 ) but requires up to 2 orders of magnitude fewer metadata accesses per batch ( Appendix D.3 ) . We compare hDTR to other heuristics inspired by recent work in our simulated evaluation ( Sec . 4 ) and discuss an even broader class of heuristics in Appendix D. Our heuristic formalization in terms of s , m , and c0 is sufficiently general to express several existing heuristics for caching and checkpointing . For example , the common LRU heuristic is “ minimize 1/s ( t ) , ” the GreedyRemat heuristic from Kumar et al . ( 2019 ) is “ minimize 1/m ( t ) , ” and the MSPS heuristic from Peng et al . ( 2020 ) is “ minimize cR ( t ) /m ( t ) ” ( where cR ( t ) sums c0 over t ’ s evicted ancestors ) . Deallocation . Deallocation policies present further tradeoffs since tensors marked as deallocated by the original program are still potential dependencies for rematerializations . In principle , DTR could simply disregard deallocations by the original program , but this would ignore potentially useful information about the deallocated tensors ( viz. , that the original program will not use them again ) . Banishing ( permanently freeing ) deallocated tensors can save memory immediately and is the only way to free constants ( which can not be evicted ) ; however , it can prevent possible future evictions since the children of a banished tensor can not be rematerialized . By contrast , evicting deallocated tensors does not prevent potential evictions , though it increases the runtime ’ s management overhead and keeps constants in memory . In the heuristics we examined , we implemented an eager eviction mechanism , which evicts a tensor as soon as all external references to it are freed . This lets DTR adhere to the garbage collection pattern of the underlying framework , preempting desirable evictions , which further reduces future runtime overhead . ( See Appendix D.2 for a comparison of deallocation policies . ) 3 FORMAL BOUNDS . Following Chen et al . ( 2016 ) , we prove a bound on DTR ’ s checkpointing overhead ( for a particular eviction heuristic ) on a linear feedforward network ofN nodes . Even without the ability to inspect the model , DTR requires only O ( N ) tensor operations under a √ N memory budget , the same bound ( up to constant factors ) as the Chen et al . ( 2016 ) static checkpointing technique and the optimal Θ ( N ) required by a memory-unconstrained algorithm . We also establish that DTR ’ s dynamic approach can not always match the overhead of static checkpointing : given N tensor operations and a memory budget of B , under any deterministic heuristic , an adversary could always construct a network where DTR would perform a factor of Ω ( N/B ) more tensor operations than a ( potentially expensive , see Jain et al . ( 2020 ) ) optimal static checkpointing algorithm . Linear Feedfoward Overhead . We assume that tensor computations dominate runtime and , as in prior work ( Griewank & Walther , 2000 ; Chen et al. , 2016 ; Binder et al. , 1997 ; Beaumont et al. , 2019b ) , that each tensor is of unit space and time cost . For the proof below , we use the heuristic he∗ , which evicts a resident tensor t with minimal |e∗ ( t ) | . Theorem 3.1 . Given an N node linear feedfoward network and a memory budget B = Ω ( √ N ) , DTR with heuristic he∗ can execute one forward and one backward pass in O ( N ) operations . Proof Sketch . During the forward pass , DTR performs exactlyN tensor operations : since each node of the linear feedforward network depends only on the previous node , no rematerialization is necessary . Our heuristic he∗ , which evicts tensors with the smallest evicted neighborhoods , ensures that the B tensors resident at the conclusion of the forward pass are evenly spaced throughout the network . In turn , these evenly spaced checkpoints ensure that DTR never has to successively rematerialize too many tensors . As the backward pass proceeds and checkpoint tensors are freed , the overhead to compute all gradients between the checkpoints k and k + 1 shrinks as log ( k ) /k2 , which sums to a constant . The full proof of Theorem 3.1 is provided in Appendix A. Adversarial Overhead . Using a simple heuristic , DTR can match the performance of static checkpointing on linear feedfoward networks despite lacking advance knowledge of the architecture . However , DTR can not always match the performance of optimal static checkpointing on an arbitrary network because it can not access or reorder the network . Theorem 3.2 . For any deterministic heuristic h , there exists an N -node network on which DTR with budget B ≤ N requires Ω ( N/B ) times more tensor computations than optimal static checkpointing . Proof Sketch . Generate an adversarial network G of B linear feedforward networks joined by a common parent tensor . Using h , schedule G ’ s operations such that , at each step of DTR , the next operation is taken from the end of an entirely evicted path through G , forcing DTR to rematerialize the entire path . DTR can thus be forced to perform at least Ω ( N2/B ) operations . By contrast , an optimal static algorithm can reorder G to compute each feedforward network sequentially , requiring only N computations . The full proof of Theorem 3.2 is provided in Appendix B. Theorems 3.1 and 3.2 illustrate how DTR ’ s performance , from optimal to poor , depends on interactions between heuristics and models . We next explore DTR design tradeoffs empirically .

Summary: Contributions: a) analyzes multiple heuristics for which tensors to evict where compute overhead of rematerialization is minimal overall, b) suggested approach is just-in-time, and thus does not require any static analysis of the network. That is, unlike prior work in this area, it covers any network type with no prior knowledge, c) offers a good formal analysis of proposed heuristic in terms of its components: staleness, memory capacity and recursive replay cost - their formalization covers previously published heuristics as well.  Experimental framework is sound. And some encouraging results are shown delivering memory capacity saving of 30% to 90% with training slowdown of 2x or less.