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: The paper presents an online algorithm for dynamic tensor rematerialization.  Theoretically, it shows the same asymptotic order on the memory budget and tensor operations as of the optimal static approach.  By simulation, it shows the performance matches optimal static checkpointing in a few models.  A PyTorch prototype is implemented, which shows benefits of reducing memory footprint and increased batch size comparing with basic PyTorch models without checkpointing.