Dynamic Tensor Rematerialization
Authors: Marisa Kirisame
Steven Lyubomirsky
Altan Haan
Jennifer Brennan
Mike He
Jared Roesch
Tianqi Chen
Zachary Tatlock
Venue: ICLR 2021 Spotlight (2021)
Keywords: Rematerialization
Memory-saving
Runtime Systems
Checkpointing

Abstract: Checkpointing enables the training of deep learning models under restricted memory budgets by freeing intermediate activations from memory and recomputing them on demand. Current checkpointing techniques statically plan these recomputations offline and assume static computation graphs. We demonstrate that a simple online algorithm can achieve comparable performance by introducing Dynamic Tensor Rematerialization (DTR), a greedy online algorithm for checkpointing that is extensible and general, is parameterized by eviction policy, and supports dynamic models. We prove that DTR can train an $N$-layer linear feedforward network on an  $\Omega(\sqrt{N})$ memory budget with only $\mathcal{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.

Published as a conference paper at ICLR 2021
DYNAMIC TENSOR REMATERIALIZATION
Marisa Kirisame," Steven Lyubomirsky; Altan Haan,' Jennifer Brennan Mike He! Jared Roesch, Tianqi Chen and Zachary Tatlock* (jerry96, sslyu, altanh, jrb, dh63, jroesch] fes. wwwwington.edeii edu, eichen@cmu.edu, ztatlockecs. washington. edu
ABSTRACT
Checkpointing enables the training of deep learning models under restricted mem- ory budgets by freeing intermediate activations from memory and recomputing them on demand. Current checkpointing techniques statically plan these recompu- tations offline and assume static computation graphs. We demonstrate that a simple online algorithm can achieve comparable performance by introducing Dynamic Tensor Rematerialization (DTR), a greedy online algorithm for checkpointing that is extensible and general, is parameterized by eviction policy, and supports dy- namic models. We prove that DTR can train an N N-layer linear feedforward network on an ((vNN memory budget with only (NV) tensor operations. DTR closely matches the performance of optimal static checkpointing in simulated experiments. We incorporate DTR prototype into PyTorch merely by interposing on tensor allocations and operator calls and collecting lightweight metadata on tensors.
INTRODUCTION
As state-of-the-an 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-intense 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 demand. Adapted from techniques in automatic differentiation (Baydin et al., 2015) Griewank & Walther [2000] [Siskind & Pearlmutter 20018), 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 | aiin 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), greedy online algorithm for heuristically checkpointing arbitrary DL models. DTT operates like tensor-level cache: it collects metadata on tensors and operators as model is trained and uses it to guide heuristics that choose which activations to free and later recompute. As runtime system, DTR can utilize dynamically gathered information (e.g., measured operator costs). Additionally, its simple, eache-like approach requires no advance knowledge of the model or application, letting immediately support arbitrarily dynamic models and applications featuring higher-order differentiation. For example, given a model with data-dependent control flow like TreeL.STM (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.
1 Paul G. Allen School of Computer Science & Engineering, University of Washington, Seattle, WA [OctoML] Seattle, WA Scchool of Computer Science, Carnegic Mellon University, Pittsburgh, PA

---

Published as a conference paper at ICLR 2021
Performop(op, args):
kematerialize(D)
Note: Perforss op(args), rematerializing any evicted arguments. wraps every operator invocation.
args := operator and arguments that
produced (from metadata)
Performop(op, args)
Exclude members of args from eviction for any evicted arg in args:
A1 Allocatesufferbb::
Remateriallzeeaar))
Note: wraps every memory allocation.
AlToaateurffrrssscc(((rssss)) i= call op(arg)) store into buf Permit eviction for members args again Update metadata for and return res
wh1le availabll memory < b:
PerformEviction()
return new buffer of size b
Deallocate(t):
Note: wraps every tensor
PerforaEviction(): the tensor chosen by the heuristic
Heuristic decides policy for t (e.g., free permanently or simply evict)
Execution Trace
Before
4
Computing with memory budget 4:
Performop(op, ts. t2)
t1
[t,, become unevictable]
Rematerlalize(tt)
t
t
(erformop(op,, t₂)
[t; becomes unevictable]
Al1ocateBuffer(t:.stt
After
t
4
C6
Performeviction() #eg, to
becomes evictable]
allocateeuffer(tt... size)
L
O
MEMORY
Perfornevictfon() eeg,
Figure I: (Top) Pseudocode for DTR's basic logic (independent of heuristic), and (Bottom) DTR's sequence of events in an operator call. Note that PerformÖp () may make further recursive calls in order to rematerialize arguments.
This paper describes DTR's design (Sec. and makes the following contributions:
We prove that DTR can train an N-layer linear feedforward network on an S2( (2(VN) memory budget with only (NN tensor operations (Sec. 3). which within 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 greedy algorithm parameterized by maching-inspires heuristics. In simulated trials our heuristic attains variety of DL models (Sec. 4
near-optimal performance
We implement DTR prototype by making only modest modifications to the PyTorch frame- work, enabling training under restricted memory budgets for both static and dynamic models and demonstrating the ease with which algorithm be incorporated into existing DL. framework (Sec. 55.
Note that techniques other than checkpointing, such as swapping tensors between devices, can also enable training under limited memory. In Sec. we discuss these approaches and how they could operate with DTR.
DYNAMIC TENSOR REMATERIALIZATION
We introduce Dynamic Tensor Rematerialization (DTR), thin runtime layer that intercepts tensor allocations, accesses, and deallocations and eliminates the need for ahead-of-time model analysis support checkpointing. Figure [ shows DTR's high-level approach. When a tensor allocation occurs wallocateBuffee.
DTR first checks if sufficient memory is available. If so. it generates fresh tensor identifier, initializes 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) cannot be evicted since no corresponding operation rematerializes them. Upon tensor access, DTR first checks the tensor is resident in memory. If so, 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: the arguments evicted tensor's parent operation have also been evicted, then they must first be
metadata for future recomputation, allocates the requested
2

---

Published as a conference paper at ICLR 2021
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 sequence of tensor operations without any inherent requirement to recognize training-spectt 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 single operation cannot fit into available memory, rematerialization will fail; therefore, on a given model and input, there may be 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 caching, DTR's eviction heuristic dynamicullyy predicts which resident tensors are least valuable. The choice of heuris- tic determines what metadata (additional runtime facts) must be tracked for each tensor and operator and thus affects DTR's 's runtime overhead. In our evaluation, we consider 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, (t), the time required compute from its parent tensor(s). We observe that DTR's s metadata overhead is low relative to the cost of typical DL tensor operations.
We propose a remaatrializzaioo-sseciif heuristic that balances staleness, memory, and cost, evicting the tensor that is stalest (least likely to be needed soon). largest (saves the most space), and cheapest (requires the least additional rematerialization if is needed again). To capture the total amount of rematerialization required if is evicted, we sum the costs over the tensor's 's evicted neighborhood e" (t), i.e., the set of evicted tensors that would either need to be rematerialized to recompute or would need t to be resident to be recomputed. We define the projected cost, c(t), of rematerializing tensor as (o(t) + "x'ee(() (t').. Using this definition, we define our heuristic, which evicts the tensor minimizing /NDTR (t) = c(t)/[m(t)- s(t]]. By including both forward and backward dependencies of f in e (t), lIDTRR 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 I where the resident tensors are (to, f2, f3, to). Before node ty is computed, we have e (tz) { t and e" (ta) = (t1, 44, ts), Since each new eviction can expand a given tensor's evicted neighborhood and each rematerialization 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 MPTR analogously (Sec. 4.1 and Appendix C.3 contain details), which performs nearly as well as hDTR in our evaluation (Sec. but requires up to orders of magnitude fewer metadata accesses per batch (Appendix D.3).
shrink it, dynamically tracking evicted neighborhoods can introduce
We compare hDTR other heuristics inspired by recent work in our simulated evaluation (Sec. and discuss an even broader class of heuristics in Appendix [ Our heuristic formalization in terms of 8, m, and co sufficiently general to express several existing heuristics for caching and checkpointing. For example, the common LRU heuristic is "minimize 1/s()), the GreedyRemat heuristic from Kumar
ct al. (2019) "minimize 1 1/m((),, and the MSPS heuristic from Peng et al. (2020) is "minimize m(()"" (where ()(() sums Co over I's evicted ancestors).
ca(1
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 dcallocations 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 cannot be evicted); however, can prevent possible future evictions since the children of a banished tensor cannot be rematerialized, By contrast, evicting deallocated tensors does not prevent potential evictions, though it increases the runtime's 's management overhead and keeps constants in memory. In the heuristics we examined, we implemented an eager eviction
3

---

Published as a conference paper at ICLR 2021
mechanism, which evicts tensor 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 dcallocation policies.)
FORMAL BOUNDS
Following Chen al. (0166), we prove bound DTR's checkpointing overhead (for particular eviction heuristic) on linear feedforward network of N nodes. Even without the ability to imspcct the model, DTR requires only ((N) tensor operations under / 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 nemory-unconstrained algorithm. We also establish that DTR's dynamic approach cannot always match the overhead of static checkpointing: given N tensor operations and memory budget of B, under any deterministic heuristic, an adversary could always construct a network where DTR would perform factor of S2(N/B) more tensor operations than (potentially expensive, see Jain et al. (2020) optimal static checkpointing algorithm.
Linear Feedfoward Overhead. We assume that tensor computations dominate runtime and, in prior work (Griewank & Walther 2000; [Chen et al.] 2016 Binder et al., [997] Beaumont et al.) 20T9b), that each tensor is of unit space and time cost. For the proof below, which evicts a resident tensor t with minimal (t)l-
the heuristic he*
Theorem 3.1. Given an N node linear feedfoward network and a memory budget B = Q(VN), DTR with heuristic he can execute one forward and one backward pass O(N) operations.
Proof Sketch. During the forward pass, DTR performs exactly N tensor operations: since each node of the linear feedforward network depends only on the previous node, no rematerialization is necessary. Our heuristic h.,, 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 & and A + 1 shrinks as log(k)/k², which sums to constant. The full proof of Theorem 3.1 is provided in Appendix A
Adversarial Overhead. Using a simple heuristic, DTR match the performance of static check- pointing on linear feedfoward networks despite lacking advance knowledge of the architecture. However, DTR cannot always match the performance of optimal static checkpointing arbitrary network because it cannot access or reorder the network.
Theorem 3.2. For any deterministic heuristic h, there exists an N-node network 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 linear feedforward networks joined by common parent tensor. Using h, schedule G's operations such that, at each step of DTR, the next operation taken from the end of an entirely evicted path through G, forcing DTR rematerialize the entire path. DTR can thus be forced to perform at least (2 /B) operations. By contrast, an optimal static algorithm can reorder 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 interac- tions between heuristics and models. We next explore DTR design tradeoff's empirically.
4 HEURISTIC EVALUATION
We simulated DTR on a variety of models to empirically evaluate its checkpointing performance across different heuristics and compare it to the static checkpointing schemes examined in Jain et al.) (2020). DTR enables training under restricted memory budgets and closely matches the performance of optimal baseline.
4

---

Published as a conference paper at ICLR 2021
4.1 HEURISTICS EXAMINED
We examine variants of the evicted neighborhood based PDTR heuristic described in Sec. 2 (on which we establish formal bounds) as well as heuristics inspired by past work in caching and checkpointing. All following heuristics are defined as a score function in terms of the metadata m(t), (t), and (t), where the tensor with the minimum score is evicted.
In addition to hDTR, we consider NDIR. which uses an equivalence class based approximation ě* for e", and h which only uses individual tensors' costs instead of costs over evicted neighborhoods. We compare against other variants of PDTR in Appendix [ but here we focus on these in particular because (1) bork lets us assess the importance of tracking evicted neighborhoods at run time, and (2)
1DTR lets us evaluate how well εx approximates e" practice. We define the HDTR variants as:
def Co(() + Eve co(')) m(t) s(t)
def co(t) + (o cu(t') m(t) s(t)
def co(() m(t) s(t)
hDIR
NDIR
DDR
Rather than using directed dependencies, (t) treats the dependency graph of tensors as undirected (thus admitting some spurious dependencies), letting decompose the graph into a of disjoint evicted components. We can track these evicted components efficiently using a union-find data structure with running sum for each component. When a tensor is evicted, its component is unioned with those of any evicted neighbors and co(t) is added to the component's 's running sum. Though this enables near-constant-tinn sums), union-find does not support splitting. To efficiently split components, we make another approximation: when tensor is rematerialized, we simply subtract (t from its component's running sum and map to new (empty) union-find component. Since this approach removes no connections, produces "phantom dependencies" between some tensors. In practice, we find that despite these additional dependences, hDTR closely matches the performance of hyлR (Figures 2 3 but requires fewer operations per eviction and rematerialization. See Appendix C.3 for a more detailed description of 2* (1).
merging between components (by unioning and adding the
We also consider the following heuristics inspired by past work:
(t) del 1 s(t)
(t) def 1
hMMPss() del Co(() + Ereen(o Co(t')
(t)
m(t)
where (R(()) is the set of evicted tensors that would have to be rematerialized in order to remate- rialize ALRU is based on the common "least-recently used" policy for caching, hyize is based on GreedyRemat from Kumar et al. (2019) (used in TensorFlow XLA), and hMsPS is based on the MSPS heuristic from et al. 2020). We also include a random baseline, hrand (t) del U 1). to assess how well a heuristic using no metadata whatsoever performs.
4.2 COMPARING DTR ACROSS HEURISTICS
Experimental Setup. To model realistic execution setting for DTR, instrumented Py- Torch (Paszke et al., (019)) to log operations performed, metadata on tensors and operators (including sizes, compute times, and parent tensors), and deallocations during the execution of various models. We replayed the logs in a simulator that models the behavior of DTR in the style shown in Figure [] The simulator tracks the tensors in memory any given time, chooses tensors evict per the heuristic when the memory budget is exceeded, and sums the total cost of the model operators and rematerial- izations. For verisimilitude, the simulator also models the semantics of various low-level Py PyTorch implementation details, including tensor aliasing, in-place mutation, and multi-output operations. We gathered logs from several static models examined in recent work, such as Jain et al. (2020) and Peng al. 2020), in addition to three dynamic models (LSTM, TreeLSTM, and Unrolled GAN); each log corresponds to an execution of the forward pass, computing the loss, and performing the backward pass. The simulator also enforces the additional condition that gradients for all trainable weights be resident at the end of the simulation in order to model the requirements for performing a full training step. Appendix C gives a full technical specification of the simulator and log format.
Results. For all models in Figure?) 2 DTR executed a training step using a small fraction of the normal memory required with limited compute overhead. Furthermore, unlike existing static approaches,
5

---

Published as a conference paper at ICLR 2021
2,00
InceptionV4
Transformer (10)
U-Neet (6)
TreeLSTH
3.50
Computo Overhead
1.25
1.00
2.00
ResNet-32 (56)
NenseNet- 121 (84)
LSTM Your dervcsion A1r
Unrolled GAN siees 5325512
3.75
224x220
22
Sunveric longth 126
2.50
1.25
1.00
0.1
0.3
0.5
07
0.9
0.1
0.3
0.5
9.7
0.9
0.1
0.3
0.5
0.7
0.1
0.3
0.5
0.7
Memory Ratio
Rpm
Paccrs
++- Prent
Figure 2: Simulated results comparing different heuristics various models, showing the rate of computational slowdown for different budgets (fractions of the original peak memory usage). The black area in each graph corresponds to the memory required to store inputs and weights, while the gray area denotes the single operator requiring the most memory to be live at once. The dashed and dotted lines represent the last ratio before thrashing (2 2x slowdown) and out-of-memool errors, respectively. All logs were produced by running each model 50 times on a single input on a machine with an NVIDIA Titan V GPU (CUDA 10.1, CuDNN 7.6.4) and 16-core AMD Ryzen Threadripper 1950X on Ubuntu 18.04, logging the final "warmed-up" run.
DTR automatically supports models with arbitrary dynamism. In all cases, results show that heuris- ties incorporating more information about chain rematerializations (hDTR, NDTR- and hMsps) can operate lower budgets and perform fewer rematerializations than heuristics using less information. However, these complex heuristics also introduce more runtime overhead, which must be considered when implementing DTR. In particular, our simulations showed that lDTTR incurred up 2 orders of magnitude more metadata accesses per batch compared to ADTR- and up to 3 orders of magnitude more compared to loccall (see Appendix D.3). The fact that MPTRR closely matches the performance of hDTR while incurring much less runtime overhead suggests that it would be more effective in practice. Note that even simple heuristics like /LUU, which require only modest runtime overhead, typically enabled training with 30% less memory.
4.3 COMPARING DTR TO STATIC TECHNIQUES
We compared the performance of DTR using lDTR, M and (as simple baseline) /L against static checkpointing techniques, including the optimal Checkmate tool of Ja et al. (2020). As Figure shows, DTR's hETR and /PTR heuristics obtain performance remarkably close to Checkmate' optimal solutions; even the much simpler ĀLRU heuristic obtains superior performance relative to the static baselines. While Checkmate requires full ahead-of time knowledge of the model and seconds or minutes per budget to compute guaranteed-optine solutions using an integer linear programming (IIP) solver, DTR finds comparable solutions dynamically and in milliseconds without ahead-of-time knowledge of the model.
5 PROTOTYPE IMPLEMENTATION
We implemented a DTR prototype in PyTorch and evaluated its performance on variety of models. We chose PyTorch because its eager mode execution ("define by run"") accomodates arbitrary control flow in models but makes static analysis more difficult; hence, it is setting where DTR's Rss online nature is an asset. Per the results in Sec, 4 we implemented hDTR as the prototype's heuristic. The core system was implemented in only 1,161 lines of code and made no deep modifications to PyTorch's 's memory management internals or tensor abstractions, illustrating the simplicity of our system. The remaining 2,647 lines of changes were primarily boilerplate operator overloads used to dispatch tensor operations through DTR's 's core logic (Appendix E. describes our prototype implementation's is structure).
Publicly available at wttpt:tt///// http:e///tthhhhoccc/////mm///ttt

---

Published as a conference paper at ICLR 2021
2.4
3.5
Eren cta.. VIT
3.4
3.4
Grizann Walther logiti
Overhesd (x)
Erecorditaal (itaa
Crepemate Joptima (IF
Pers
1/2
Poo
3/1
1.1
1.1
Budget (GB)
Figure 3: DTR's overhead from operators is competitive with Checkmate's, which uses ILP to produce optimal rematerialization schedule. This comparison extends Figure 5 in Jain et al. (2020) by adding the DTR simulator as a "solver" that translates Checkmate's 's Keras-basee based graph representation into the DTR simulator's representation. To produce this comparison, Jain et al. 2020))s 's evaluation artifact because the PyTorch logs from Sec. 4.1 did not contain some information that past checkpointing techniques require (such which backward operators correspond to which forward ones). Also included in the comparison (from the original experiment) the Griewank & Walther (2000) Treeverse algorithm and variants of the Chen et al. (2016) checkpointing algorithm (modilied to handle skip connections like those in ResNet).
modified
LSTM (28)
Bimraion sion. 512,
InceptionV4 (84)
Unrolled Gan
Transformer 512 (Bqqaaax Lenggit)
Socusnne Longgh 128
66 Slups, 512x S12
5001
Time (ms) / Batch
250
2001
11.1
10.0
6.0
4.0
11.3
10.0
8.0
6.0
4,0"
11.2
10.0
8.0
8.0
6.0
UNet:::
ReaNet1 202 (84)
TreeLSTM (9)
DenseNet121 (84)
616% 003
224 224
Bireey uourwh degalt mode slse 640
224
221
10:0
300
500
200
8.4
8.0
7.0
6.0
4.0"
8.9
8.0
6.0
4.0
2.0
8.9
8.0
6.0
4.0
2.0
11.0
10.0
8.0
6.0
4.0"
Memory Budget (GIB)
Operttor
Compute
Ediition
Rematofolttaton
Urprofiled
Pyforrhh
Figure 4: We profiled the running time of our prototype for various models and memory budgets on machine with an NVIDIA Titan V GPU (CUDA 10.1, CuDNN 7.6.4) and a 16-core AMD Ryzen Threadripper 1950X on Ubuntu 18.04. The red dotted lines correspond to trials that either ran of memory or thrashed (2 2x unmodified PyTorch's time). Model batch sizes are given in parentheses. To ensure the accuracy of the DTR prototype's profiling, we used PyTorch's ' synchronous computation mode (see Appendix [E.T]. Results (mean of 100 trials) are compared against unmodified PyTorc.. "Cost compute" (computing heuristic scores) and "eviction loop" (comparing scores over tensors) correspond to overhead from DTR runtime itself, which can be reduced by more efficient implementation. "Unprofiled time" is the remainder of the time per batch; may be due to runtime overhead from parts Py Torch not modified in the prototype, like the operator dispatch system. The large proportion of unprofiled time in Unrolled GAN is likely due its extensive use of Python reflection. The budgets with asterisks were run with the random sampling optimization (see Appendix E.2) disabled, as sampling caused occasional failures those budgets.
Our empirical evaluation demonstrates that DTR can efficiently train models under restricted memory budgets using the AoT heuristic. We used the same models and experimental setup in Section 4 timing the forward pass, loss computation. and backward pass. Table [bbleepeeenett several cases where DTR trains models much larger input sizes than unmodified PyTorch, including dynamic model,

---

Published as a conference paper at ICLR 2021
ResNet-1202 (Batch Size)
Transformer (Batch Size)
UNet (Batch Size)
TreeLSTM (Tree Nodes)
100
120
90
2
2
2
2
DTR
0.974s
1.18s
1.285
1.39
367ms
830ms
950ms
1078ms*
566ms
684ms 822ms*
1370ms*
0.486s
1.05s
2.50s
7,8%=
PT
0.712s
X
X
X
331ms
X
X
481ms
X
x
0.4315
X
X x
Table 1: Median execution times per batch (out of 100 runs) for various models, giving both the largest input size that unmodified PyTorch ("PTT) could support on our GPU and larger input sizes DTR could support. Input sizes are as in Figure except for TreeLSTM (complete binary trees with nodes of size 1024 x 1024) and Transformer (sequence length 256). Asterisks indicate inputs on which the random sampling optimization was disabled due to occasional failed trials. Even without sampling, DTR still occasionally failed on UNet (see Appendix E..3 for details). This behavior may be due to PyTTrrch memory allocator implementation details or poor rematerialization decisions influenced by variance in individual operator times.
TreMLSTM. This highlights that DTR enables exploration of models that push the boundaries of existing deep learning architectures. While the simulated trials in Sec. 4.2 consider the slowdown due only to rematerializations but not overhead from managing metadata and computing heuristics, Figure measures the time per batch required to train eight DL models on a variety of restricted memory budgets. profiling the time spent by the runtime system. Among the models is Unrolled GAN, which uses higher-order partial derivatives and Python reflection extensively; the DTR prototype supported these unusual features, underscoring its generality. Despite prototype's simplicity it merely loops through all tensors when searching for an eviction candidate and recomputes the heuristic scores from scratch each time on most models, its overhead due to searching and computing heuristics remains low for most memory budgets. In Appendix [E.2 we discuss two approximate optimizations included in the prototype reduce the overhead of searching over tensors and additional ways to reduce DTR's s runtime overhead.
RELATED WORK
Checkpointing in Reverse-mode Automatic Differentation (AD). Checkpointing in DL takes in- spiration from checkpointing in reverse-mode AD (Baydin et al. 2015). The latter reduce the number of values stored in the "tape" by recomputing segments of the tape (demarcated by "checkpoints"). Treeverse (Griewank 1994; Griwwnkk & Walther 1997 2000) uses a binomial partitioning scheme to mark checkpoints, achieving logarithmic growth in space exchange for logarithmic grown in computation. Later works, such as Hascoet & Pascual (2013) and Siskind & Pearlmutter (2018), extend Treeverse's 's approach to handle arbitrary control flow by inserting code at compile time to mark checkpoints according to policies (e.g.. "'checppiint every k iterations" for statically unbounded loop). Unlike DTR, these techniques do not use dynamically gathered information.
Checkpointing in DL. Many DL models can be represented as static dataflow graphs, enabling the straightforward application of Treeverse-like partitioning approaches. [Cen et al. (2016) apply this approach by dividing the network into segments to be recomputed during backpropagation, presenting schemes that allow for training an N --layeer feedforward network in ((VN) memory in Olog N) memory with (No N) additional tensor operations. Grrssyy et al. (2016) present a similar segmenting approach for recurrent neural networks, thereby supporting some dynamism beyond static computation graphs. Other recent work rematerializes individual activations rather than entire segments, attaining better bounds than Chen et al. 6) Kusumoto et al. (2019), ). Kumar et al. (2019), and [eaumont al. (2019a) apply thepheth analyses to make rematerialization plans, while aii et al. (2020) apply integer linear programming (ILP) to find optimal solutions.
with one extra forward pass (O(N) tensor operations)
DTR differs fundamentally from those approaches because it handles arbitrary dynamic control flow in models (making no assumptions about the model's structure) and operates online, giving it access to dynamically gathered information. In principle, a static checkpointing technique could be applied to a dynamic model "just in time" by unrolling the model on the fly, but some static analyses (like an ILP solver) can be too expensive to run each epoch. Unlike static approaches, however, dynamic planning introduces overhead at run time, which limits the analyses that DTR's 's heuristics can feasibly perform. Note that the Chen et al. (2016) greedy scheme and the GreedyRemat baseline in Kumar et al. (2019) are similar to DTR in that they greedily place checkpoints using a heuristic (albeit statically). However, their heuristics only the sizes of tensors.
8

---

Published as a conference paper at ICLR 2021
DL Memory Managers. Other work has enable the training of DL models lower memory budgets by swapping tensors between GPUs or to host RAM. Huang et al. (2020) use a genetic algorithm to plan swaps between devices static computation graphs. Capuchin by Peng et al. (2020) and Superneurons by Wang et al. (2018), like DTR, use runtime systems and incorporate checkpointing as well. Capuchin's is checkpointing phase, which resembles DTR's, uses dynamically gathered information for checkpointing: it performs a single batch without checkpointing (only swapping) and uses the costs measures determine where set checkpoints. However, Capuchin's 's and Superneurons's 's checkpointing schemes assume static model architecture (inferred from an initial profiling batch), which they use to plan recomputations in advance. Swapping systems like Capuchin rely on interleaving communication and computation at low level for performance, which may be difficult to apply in an online setting.
These works highlight that swapping and rematerialization are complementary approaches, raising the question of whether DTR can be combined with swapping without disrupting existing methods' overlapping of computation and communication. One possibility would be to assume a fixed swapping schedule and use DTR to replace the rematerialization schemes used by systems like Capuchin (perhaps given constraint like treating values to be swapped out unevictable). Another intriguing possibility would be to use swapping as a form of "eviction" in DTR, where the "cost" for swapped-out values would be the communication time. Swapping presents interesting tradeoffs with rematerializations since it may scale better than some tensor operators. However, incorporating swapping into DTR's online approach presents the problem of efficiently overlapping computation and communication since the runtime would need to guarantee that a computation scheduled concurrently with a swap would not need swap values back in. This could greatly complicate planning (e.g.. requiring some lookahead to avoid missed swapping opportunities) and would be fertile ground for future work.
HemmoryEEffiieet DL Model Designs. Some recent work manually modifies DL models to perform similar computations using less memory, which may be used alongside checkpointing and swapping approaches. One example is the use of reversible layers, which enable recomputing a forward value during backpropagation using the result of the following layer. [oomzz et al. (2017) and Kitaev et al.) (2020) employ reversible layers to create versions of ResNet and Transformer, respectively, that can train using less memory.
7 CONCLUSION
DTR provides a simple, customizable approach to checkpointing for DL models. It supports a broad range of applications without the need for any ahead-of-time analyses, manual annotations, or modifications. Our formal results establish that DTR can match the same asymptotic bounds as recent static checkpointing approaches for linear feedforward networks. In simulation, it enables training for a range of both static and dynamic models under various restricted memory budgets and closely matches the performance of optimal checkpointing. The DTR prototype in Py PyTocc demonstrates how our approach can be incorporated into existing frameworks with modest, non-invasive changes by simply interposing on tensor allocations and operator calls and collecting lightweight metadata on tensors. Our results also open several avenues for future work. For example, DTR could easily be extended to leverage additional information that may further reduce runtime overhead, such as learning from past batches.
ACKNOWLEDGEMENTS
This work was supported by the Applications Driving Architectures (ADA) Research Center, a JUMP Center co-sponsored by SRC and DARPA. The Titan V used for this research was donated by the NVIDIA Corporation. We thank Paras Jain and Aniruddha Nrusimha for assistance in setting up and running the Checkmate MLSys 2020 artifact and providing helpful additional information about the Checkmate tool. We are grateful to Edward Z. Yang for helpful advice modifying PyTorch. We acknowledge Yonghao Zhuang for drawing our attention to an omission in our description of the hDTR splitting approximation in Section 4.1 which have corrected in this version. We also thank Sandy Kaplan, Eunice Jun, Josh M. Pollock, Samuel Ainsworth, and Sam Kaufman for providing feedback and useful comments various drafts of this work.

---

Published as a conference paper at ICLR 2021
REFERENCES
Atilim Gunes Baydin, Barak A. Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic differentiation in machine learning: survey. CoRR, abs/1502.05767, 2015. URL
http://arxiv org/ abs///b////// 1502.500.000 05767
Olivier Beaumont, Lionel Eyraud- Julien Herrmann, Alexis Joly, and Alena Shilova. Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory. Research Report RR- RR-9302, Inria Bordeaux Sud-Ouest, November 2019a. httph:t/////// //hal
Anria.frhhae -
Olivier Beaumont, Julien Herrmann, Guillaume Pallez, and Alena Shilova. Optimal memory- aware backpropagation of deep join networks. Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences, 378, 01 2019b. doi: 10. (0098//sta.201..0000
John Binder, Kevin Murphy, and Stuart Russell. Space-efficient inference in dynamic probabilistic networks. Proceedings of the Fifteenth International Joint Conference on Artifical Intelligence Volume 2, IJCAI'97, pp. 1292- 1296, San Francisco, CA, USA, 1997. Morgan Kaufmann Publishers Inc. ISBN 15558604804.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. CoRR, abs/1604.06174, 2016. http::/aaxx...oorr/t
06174
Aidan N Gomez, Mengye Ren, Raquel Urtasun, and Roger B Grosse. The reversible residual network: Backpropagation without storing activations. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances Neural Information Processing Systems 30, pp. 2214 2224. Curran Associates, Inc., 2017.
Andreas Griewank. Achieving logarithmic growth of temporal and spatial complexity re- verse automatic differentiation. Optimization Methods and Software, 1, 04 1994. doi:
10.1080/10557789000555555
Andreas Griewank and Andrea Walther. Treeverse: An implementation of checkpointing for the reverse or adjoint mode of computational differentiation. ACM Trans. Math. Software, 26:200-0, O, 1997.
Andreas Griewank and Andrea Walther. Algorithm 799: Revolve: An implementation of checkpoint for the reverse or adjoint mode of computational differentiation. ACM Transactions on Mathemat- ical Software, 26 1996:9999445. 45, mar 2000. ISSN 0098-3500) URL fattp://ce //oo1. acm. org
1145/347833 347.4447844 347846
Audrunas Gruslys, Rémi Munos, Ivo Danihelka, Mare Lanetot, and Alex Graves. Memory-efficient backpropagation through time. CoRR, abs/1606.03401,, 2016. URL htttp://axxiv.ort
0is/16666.04001
Laurent Hascoet and Valérie Pascual. The tapenade automatic differentiation tool: Principles, model, and specification. ACM Trans. Math. Sofnw., 39(3), May 2013. ISSN 0098-3500, doi: 10. 10.1144/2450153.245500
httptts::///// //ooi. orgiilorg/lle 10. 1145/14551111 2450158
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-Fifil International Con- ference on Architectural Support for Programming Languages and Operating Systems, ASP- LOS '20, pp. 1341-1355, New York, NY, USA, 2020. Association for Computing Machinery. doi: 10.1 022114//377376633888555
ISBN 0781450371025.
URIhhttps:///oi http:////o...0/10..11/00000
3373376. 33786.337733
Paras Jain, Ajay Jain, Aniruddha Nrusimha, Amir Gholami, Pieter Abbeel, Joseph Gonzalez, Kurt Keutzer, and lon Stoica. Checkmate: Breaking the memory wall with optimal tensor rematerializa- tion. In Proceedings of Machine Learning and Systems 2020, pp. 497-511, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020. URLtttps: /ttttps://opeeeeviee
10

---

Published as a conference paper at ICLR 2021
Ravi Kumar, Manish Purohit, Zoya Svitkina, Erik Vee, and Joshua Wang. Efficient rematerialization for deep networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. Alché- Alché-Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran As- sociates, Inc., 2019. URL https: / pttttss:/prrcceeednggg neur neurips......... scc/paper/200 9/file/ pdf
ffe10334251c0
Mitsuru Kusumoto, Takuya Inoue, Gentaro Watanabe, Takuya Akiba, and Masanori Koyama. A graph theoretic framework of recomputation algorithms for memory-effiiiet backpropaga- tion. In H. Wallach, Larochelle, A. Beygelzimer, F. d' Alché Buc, Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL
www.ipe. http://s....uursss.o/p/p////////
pdf
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Vetio,, Zeming Lin, Alhan Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017. Rtr.aattsss/ppeeeeeeessseeeeeder
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach De Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style. nigh-pefformace deep learning library, 2019. URL /ttp::://arxatr
01703
Xuan Peng, Xuanhua Shi, Hulin Dai, Hai Jin, Weiliang Ma, Qian Xiong, Fan Yang. and Xuehai Qian. Capuchin: Tensor-based gpu memory management for deep learning. In Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems, ASPLOS '20. pp. 891-905, New York, NY, USA, 2020. Association for Computing Machinery. ISBN 9781450371025. doi: 10.11455333336.33375506 URL
Nadav Rotem, Jordan Fix, Saleem Abdulrasool, Summer Deng. Roman Dzhabarov, James Hegeman, Roman Levenstein, Bert Maher, Satish Nadathur, Jakob Olesen, Jongsoo Park, Artem Rakhov, and Misha Smelyanskiy. Glow: Graph lowering compiler techniques for neural networks. CoRR, abs/1 000.009050 2018.
Jeffrey Mark Siskind and Barak A. Pearlmutter. Divide-and-conquer checkpointing for arbitrary programs with no user annotation. Optimization Methods and Software, 1288-33)1288-1330 Sep 2018. ISSN 1029-4937. doi: 10.1080//055678.. 2018. 1459921..
URL http://dtttttt. http://xxdddooooooo/oo/
10. 15.1100011555707 211...2..8.....1222. 1459621
Nimit Sharad Sohoni, Christopher Richard Aberger, Megan Leszezynski, Jian Zhang, and Christopher Ré. Low-memory neural network training: A technical report. CoRR, 05//1944.106311 2019. URL
10631
Kai Sheng Tai, Richard Socher, and Christopher D. Manning. Improved semantic representations from tree-structured long short-term memory networks. Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference URL
Natural Language Processing (Volume 1: Long Papers), 2015. doi:
http: /ttp://xxxcoomm doi. wtt../...//../////.... 3115/v1/p15
1150
Linnan Wang, Jinmian Ye, Yiyang Zhao, Wei Wu, Ang Li, Shuaiwen Leon Song, Zenglin Xu, and Tim Kraska. Superneurons. Proceedings of the 23rd ACM SIGPLAN Symposium on Principles and Practice of Parallel Progrumming, Feb 2018. doi: 10.1145//178448.3777999.. URL h
7/dx doi. .asil.rg//le 10. 1145/317841111 11445/118848833188449
11

---

Published as a conference paper at ICLR 2021
A PROOF OF THEOREM 3.1
In this section, we provide proof of the ((N) runtime of DTR linear feed-forward network with uniform operator compute and memory cost, under a reduced heuristic. We begin with a thorough treatment of the network architecture, and then motivate our reduced heuristic he in this simplified setting. Finally, we prove Theorem 3.1
A.1 NETWORK DEFINITION
We assume the network consists of operators fi...., JN. where the tensor computed by the ith operator is given by ft),; with ł, denoting the tensor computed by the jth operator. Note that we consider to to be the input tensor, which for simplicity will always reside in memory and not contribute to the active memory consumption. For this reason, we may consider (1 to nullary operator. Additionally, we assume that the size of each tensor (denoted m(()) is 1, and likewise for the compute time fi( for each operator ft. Note that we may write co(ft) to mean the same as (ff) for t = t, when the index i is not convenient.
For backpropagation, assume each operator has associated gradient operator fis which computes the result fi(,-1, (4+1. We may consider fN+1 1 be an unevictable unit tensor, as the case automatic differentiation, but for simplicity we define = f1((2) and iN = (fN--11). As above, we assume unit memory and compute for each ft.
t
t2
f
N2
IN-1
tN
i
i2
is
{N-2
iN-1
in
A.2 LIVENESS AND BANISHING
To optimize memory usage during computation, we introduce the notion of liveness and banishing. At high level, liveness allows us to determine when given tensor is no longer required for subsequent network computations, which in turn allows us to permanently free (banish) tensors to regain memory when certain conditions are met.
To be more precise, we formalize the network as program:
let t = f((; let to == fz(t1);
let In == [N((N-1);
// Backpropagate.
let iN == fN(tN-1)i
let fx- = (N-1(tN-2,
let f = i2()
let & == f((tz);
We say tensor is live when there is a pending operation in the program that takes as an input. When is no longer live, and every tensor directly computed using is in memory banished, then we say is banished and we reclaim the memory used by t. Banishing a tensor additionally makes its children unevictable.
12

---

Published as a conference paper at ICLR 2021
Thus for example, ty can be immediately banished after computing, ---1 be banished after EN. both 12-2 and Ex after -1- and so on. This will become important in the proof.
The analysis of liveness can be done statically for static models, and by reference counting for models with dynamism. In both cases, liveness information is fed to DTR online through deallocation events.
A.3 HEURISTIC DEFINITION
Heuristic he is reduced form of the DTR heuristic, does not account for tensor staleness. Here, we provide a detailed motivation of its definition.
Recall the evicted neighborhood (t) of tensor as described in Section 2 and further formalized in Appendix
Definition A.1 (Projected Cost). For given tensor 1, the projected cost of i is the value
e(t)
Σ
co(fe)
(1Ee"(1)
Now, define the reduced heuristic in full generality; the definition of he will be consequence of the simplified setting we analyze.
Definition A.2 resident tensor / is defined
Heuristic (general)). The compute-memory heuristic score for a
he (t) e(t) + co(fe)
mll)
Corollary A.1. Under our simplified compute and memory constraints, he-(t) - (t) Since the heuristic is only used to rank tensors, the common additive constant is unimportant. The heuristic (e (t) will have the same behavior as |e* (t) 11
Note importantly that uncomputed tensors are not considered in any of the above definitions (as we do not know about their existence yet, from a dynamic execution perspective).
A.4 PROOF OF THROREM
Now prove Theorem 3.1 which bounds the overhead of DTR on linear feedforward network with N nodes and VN memory by a constant factor of the runtime required by algorithm with unlimited memory.
Proof. To prove this claim, we will consider the forward pass and the backward pass separately. In the forward pass, we show that our algorithm only performs N computations, matching that of an algorithm with unlimited memory. Furthermore, upon completion of the forward pass, we tightly characterize the B tensors that remain in memory. We show that a set of evenly spaced checkpoint tensors remain in memory throughout the backward pass, until banishment. The presence of these checkpoint tensors allows us to argue that the algorithm never has to rematerialize too many tensors in row. Furthermore, as the algorithm computes additional gradients, it banishes checkpoint tensors that are no longer needed, freeing more space for additional checkpoints. The overhead incurred by the algorithm can therefore be kept to constant factor of the required ((N) time. This checkpointing behavior can be seen in the trace of the algorithm, visualized Figure
We now analyze each of the phases in detail.
PILASE I: FORWARD PASS
Recall that a feed-forward network, every computation depends only on the preceding one. Thus in our simplified network, we only ever need B units of memory to compute the forward pass without any rematerializations (furthermore, this is the minimum required memory). For this reason, the forward pass requires N computations.
13

---

Published as a conference paper at ICLR 2021
Reduced CCoputeeMemoryy 2 Memory (n 200 Layers)
192
begin backprop
184
176
168
160
152
120
104
0.8
96
88
80
0.6
72
64
56
0.4
48
0.2
8
0.0
66
88
110
132
151
176
198
220
242
264
286
308
330
352
374
395
413
440
505
528
550
572
594
Figure 5: Visualization of the state of memory for DTR with N 200, B = V and heuristic her A value of 0 (black) indicates the tensor is evicted banished, (red) indicates the tensor is forward value in memory, and 1.5 (white) denotes an in- in-memory gradient tensor corresponding to the forward tensor. The backward pass begins at the red vertical line; note the presence of evenly spaced checkpoint tensors (red horizontal lines) that persist in memory throughout the backward pass. Note also the recursive checkpointing behavior visible the completely red triangles of the later gaps, when there is enough free memory to avoid repeated rematerialization altogether.
early gaps the backward pass, and finally
After completing the forward pass, we can tightly characterize the tensors remaining in memory. In particular, Lemma A.J[tell us that the maximum gap between resident tensors is bounded by
2(N 2) L -
B 3-1
We note that this bound is tight in an asymptotic sense: we can keep B tensors in memory, and the forward pass is of length N, then the maximum gap must be at least N/B.
Next, we will analyze the backward pass. Key to this analysis is the claim that "not too many" of the tensors in memory at the beginning of the forward pass evieted before banishment during the backward pass. The existence of these "checkpoint tensors" allows us to argue that we do not do too much rematerialization work.
PHASE 2: BACKWARD PASS
During the backward pass, our algorithm computes gradients L. Each gradient computation relies on two inputs: 441 and 4--1- We show that neither input incurs too much rematerialization cost 44+1 because it is pinned in memory, and t;-1 because the paths evicted tensors are not "too long." " The first condition follows from the fact that t, banished after computing f,41. therefore forcing f,41 to remain in memory until it is banished. The second condition is formalized in the following lemma, proved later in this section.
Lemma A.1 (Checkpointing)) Consider execution of the DTR algorithm with B units of memory and heuristic he, applied to the graph described in section Let S be the set of tensors in memory after computation oflx ix in the forward pass. Then, C C S is of "checkpoint" tensors from the forward pass with the following properties:
1. During the backward pass, each stays in memory until it banished.
14

---

Published as a conference paper at ICLR 2021
2. The gap between neighboring tensors in C satisfies
4(N 2) LS
-
These |C] checkpoint tensors divide the forward tensors into [C] groups, indexed by k. each of 4(N-2) The total computational cost of the backward pass equal to the sum of the computational cost for each group,
length LE I B-1
CI C Σ Ck.
The second key insight in the analysis of the backward pass is that, for every group that is processed, the algorithm banishes a checkpoint tensor с E and receives a unit of extra memory. particular, at the start of processing group |C] k, the algorithm has 2 le pieces of extra memory (two from banishing the most recently used gradient and forward tensor, and & from the banished checkpoint tensors). We leverage this extra memory to process the gradients in later groups with less rematerialization overhead, using the extra units of memory to create intermediate checkpoint tensors. The following lemma describes how the cost of computing all the gradients in a group decreases as we free more memory.
Lemma A.2. Suppose we have 2 3- pieces free memory to compute all of the gradients associated with an evicted forward tensor path of length Lę- Then the number of rematerializattons needed to compute all the gradients is of order
L 12
C== O Lx
log
Applying this lemma, the total cost of the backward pass becomes
c
c
C₁
[C]
L
Lb
+
log k
k2
[C]
[C]
log h:
L +
L
1=1
k=1
|C|
4(N
2)
log k
(4(N (4(N-2) -2) B-1 +
|C|
+
2
k=1
NN
N
[C]
log k
|C|
B
B
1:2
where S hides constant factors. Note that [C] S B, since E S where S' is the set of tensors in memory at the end of the forward pass, Also note that log is a convergent sequence, so its partial sums are bounded. Therefore, we can simplify the bound to
N
CSN +
B²
Since B ( (VN), conclude that the total cost of the backward pass is O(N). Adding this to the O(N) cost of the forward pass, we see the total compute is O(N), as desired.
15

---

Published as a conference paper at ICLR 2021
A.5 PROOFS INTERMEDIATE RESULTS
Here, we present intermediate results that we used in the proof of our main result.
Lemma A.3. Consider the DTR algorithm operating with heuristic her. Suppose we seek to (re materialiie forward tensor tk for k S N, where the resident tensor preceding ty is denoted by to (with j < k). Suppose also that t, is not evicted during the computation of tx. Then, if the algorithu begins with tj in memory and with M units of memory. and runs until computing tg. then the maximum length I. of any evicted sequence of tensors between and iy is bounded by
L - -(1)))))
Proof. Proof by induction. We will show that, when the algorithm computes tensor j + i, for the maximum length of evicted sequence of tensors between ty and tj+s
i = 1.2,..., - satisfies
L, s 2 1)M(M-MMMM -
Base case. When = 1, both ł and fy+1 = t are resident tensors, so the gap is L1 =
Inductive step. Consider the contents of memory after computing ty+s. We begin by partitioning tensors ty..... t++s into segments Si...., SM. each ending resident tensor (note, the last segment must end resident tensor, since fi+z was just computed). If < M so that there not M resident tensors, then the length of each segment zero and we are done. Otherwise, each segment corresponds an evicted sequence of zero or more tensors (i.e., the tensors preceding the resident tensor). Let denote the resident tensor that ends segment i.
Now, consider all adjacent pairs of segments (S,, S141) for s M 1. The average length of the pairs is given by
Si + S1111
L
M
M
|S|
| + [Sar
2
M
M --1
2
M
||| | + [ M 1
Σ
S
M
1
1=1
2i
S1 + |SMl
M 1
M
2(i
M
Let (Sy, Sy+1) be the pair of adjacent segments with minimum combined length. Since the average length is bounded by the inequality above, follows that the length of (Sp, Sy41) is also less than or equal to 2(i 1)/(M 1).
Since the heuristic evicts the tensor that results in the smallest gap, we conclude that the eviction will create gap larger than 2(i 1)/(M 1). By the inductive hypothesis, the largest previous gap was no larger than 2(i 2)/(M M - 1), so we conclude that the largest gap after this computation is no more than 2(i 1)/(M) - 1).
PROOF LEMMA
A.
Proof. We will prove this lemma by dividing the backward pass into two phases. In the first phase, the first two gradient computations of the backward pass, may be forced evict some element of S. In the absence of further information on the evicted tensor, we upper bound the resulting gap by twice the maximum gap between tensors in S. This gives the upper bound in Item 2 of the lemma.
16

---

Published as a conference paper at ICLR 2021
In the second phase, the remaining N 2 gradient computations of the backward pass, show that heuristic he- never leads us to evict a tensor that would lead to a gap of more than 1(N-2) tensors in memory. This allows us to conclude that the checkpoint tensors C. remain in memory until eviction, as claimed.
among the
8-1
We now elaborate on the two phases, as discussed above.
Phase 1: The first two gradient computations of the backward pass.
We present a detailed treatment of the first two gradient computations in the backward pass, iN and -1-1- We will show that, during the course of these two computations, at most one tensor from S is evicted from memory. Since Lemma A.3 tells us that the maximum gap in S satisfies L. S -(N-2) B-I we conclude that removing single tensor results in gap in C of more than 2Ls. Additionally, we will show that after the computation of the first two gradients, there are at least two non-checkpoint tensors in memory. Since only two free units of memory are required to rematerialize path of tensors, this sets us up for the analysis of the remaining gradient computations.
We begin by noting that, after the forward pass completes, tv and - both in memory (since by has just been computed, which requires 11- Since fx is no longer needed in subsequent computations, it is immediately banished. Assuming B N, this leaves with exactly one unit of free memory (if B > N, no elements of S are banished in the first two computations, and the 2L, bound is trivial). This single unit of memory is then filled by the computation of iN, which only depends on -1
Now, fN-1 is no longer needed, so it is banished, and we have exactly one unit of free memory. To compute f we require N2 and iN be in memory. Since in was just computed, it is clearly in memory. However, IN may or may not be in memory. We consider the two cases separately.
If yy- is in memory, then we immediately compute f Next, tensors y and f are banished, leaving us with the desired two free units of memory.
If, on the other hand, f is not in memory, we must rematerialize it. Let tj be the resident tensor that terminates the evieted path of tensors containing x We need to perform the sequence of computations (ty+1, ..,+,....,,,.22) However, we only have one unit free memory, so after computing 1i+1 we will need evict some tensor from memory. The evicted tensor must be 4; for some S j. as neither tj+1 nor fn can be evicted (the former will be used for the next computation, and the latter is pinned in memory).
Regardless of which tensor & is evicted, the length of the evicted path it creates cannot exceed 2Ls. where Lg is the length of the longest path in S. Lemma A.3 bounds Ls < B-1 algorithm maintains Item of the lemma.
2 so this step of the
It remains to show that the maximum gap in C does not become larger than 2Ls during the remaining steps of rematerialization, and that the computation of iN-1 ends with at least two units of free memory. To show the first claim, we note that the number of evicted tensors on the path to Ê does not exceed 2T.g (this is the maximum length possible, if f, was evicted and its adjacent evicted paths were both of length Ls). Therefore, when performing the intermediate rematerializations necessary to rematerialize 22, it is always possible to evict tensor between ty and i 2. with heuristic value of less than 2Ls. Since we evict the tensor with the smallest heuristic value, we will never create evicted path of length greater than 2
Finally, we note that, after computing (n-1, both 22-2 and iN will be banished. This leaves us with the desired two units of free memory.
We have shown that, after computing ix-1, the algorithm has two units of free memory, and the checkpoint set C' has maximum gap of no more than 2g. Next, we show that this set C' is maintained throughout the remainder of the backward pass.
Phase 2: The remaining N - gradient computations.
The analysis for the remainder of the backward pass follows via induction, using the argument for rematerializing -2-22 above.
17

---

Published as a conference paper at ICLR 2021
We have already shown base case; f2-2. For the inductive step, consider the computation of & for < - 1. Suppose we have at least two units of free memory, and f in memory. Furthermore, suppose that the set C satisfies the properties of the lemma. We need rematerialize f;-1, which terminates a path of evicted tensors of length no more than 2 As we rematerialize this path, it may require evicting tensors from memory. However, by the same logic evict a tensor resulting in a path of less than 2Ls. The algorithm will always choose this option in favor of creating a longer evicted path. We conclude that the upper bound of 2 is preserved when computing 4. Furthermore, after I; is computed, we may evict 4++1 and 4 giving us two units of free memory. This proves the inductive step.
maintain the desired properties of when computing
applied above, know that the algorithm may always choose to
Note that, in the case that = 1, the computation requires no rematerializations, as & only depends on fg, and the latter memory at the time of computing é1.
PROOF OF LEMMA A.2
Proof. Let Ci. denote the cost of processing gradient in this group. Since there are L associated gradients, the total cost is
L
C
C
To compute each Ci,, we note that computation the gradients proceeds in phases. When the first gradient is computed cost Ca,s = L&). two units of memory must be devoted the current tensor computation, while the remaining & units of memory are used for intermediate rematerialized tensors. Applying the intermediate checkpointing lemma. A.4 we conclude that some of these intermediate tensors will remain as checkpoints (indexed by j. with j = 1 indicating the highest-indexee tensor), with adjacent checkpoints separated by distance at most Lk.j 4(L, 2) We express the total cost of computing the gradients in this gap
C = Lsk +
Σ Σ Ci.k J iEgroup
We begin by considering the first group be processed, j associated with the last path between checkpoints. Since it is the first group to be processed, it has no spare memory for intermediate checkpoints. Therefore, computing the first gradient requires rematerializing the entire group (with at most LRg intermediate tensors), computing the second gradient requires rematerializing at most Lk 1 tensors, and so on. This gives total cost bounded follows (using Z denote inequality up to constant factors).
Lass
Σ iEgrap
Cl
S
Lkj
(Ln.))
4(Lx -2)
+1
k-
24
E
Next, compute the total cost of calculating all the gradients between checkpoints and 1. When the algorithm begins to compute group j. it has j pieces of extra memory, allowing it to further subdivide group into j + intervals. By the intermediate checkpointing lemma, each of these
18

---

Published as a conference paper at ICLR 2021
4(Lx,) -2 5-1
intervals is of length most
+ 1. We have
2che -2)
4(Ln.) 2)
Σ iEgroup
Ca S
j-1
1=0
[4(Lk.) -2)
Jj
+1)
j
k.j. j
Summing over the at most & checkpoints j, we conclude
Cs s Lx Σ
j
La + L Hs
L
< La log k
k2
where Hg is the ht harmonic number.
Lemma A.4 (Intermediate Checkpointing). Consider the behavior of the DTR algorithm using the heuristic her, when computing gradients for the backward pass. Suppose, immediately prior the computation of gradient tú we have + k pieces of free memory (k 2 0), and that f,41 is in memory. Suppose also that forward tensor t, is the first resident ancestor of i, so that we will rematerialize 4--1 starting from lj to compute i,, Finally, suppose that lj is never evicted until it is banished.
Then, immediately after computing iú memory contains a set of "checkpoint" tensors C with the following properties:
1. The tensors in C remain in memory until they are banished.
2. The gap between neighboring tensors satisfies
2(( - 1) LS
k +1
Proof. We begin by analyzing the state memory after computing f, Since we started with 2 2 pieces of free memory, and rematerialized 14-1 starting from lj. Lemma A.3 tells us that, after rematerializing t,-1, the gaps in memory between ty and t are all bounded by
2((i 5) 1) I.
+
We need evict additional item from memory, in order to compute é, After this single eviction, the maximum gap is no more than doubled. We conclude that, after computing the first gradient, the maximum gap is more than 2L.
remains to show that the maximum gap in does not become larger than 2L during the remaining steps of rematerialization. To show this. first note that the computation of the next gradient, f--1, begins with two units of free memory (having just banished f and t.). We also note that the number of evicted tensors that need to be rematcrialized for this gradient computation does not exceed 2L. Therefore, when performing the intermediate rematerialization necessary to rematerialize f;-2, is always possible to evict a tensor with heuristic value less than 2L. Since we evict the tensor with the smallest heuristic value, will never create evicted path of length greater than 2L.
This argument can be applied for every gradient computed between f, and f which shows that the desired properties of C are maintained.
19

---

Published as a conference paper at ICLR 2021
Figure 6: An example construction of adversarial graph. Gray tensors in memory (to must always be in memory). The initial tensor to has B paths descending from it, so there is always some path from to with no resident tensors. The adversarial construction chooses place the next node the end of such an entirely evicted path.
B PROOF OF THEOREM
In this section, we provide proof of Theorem 3.2 which lower bounds the number of tensor compu- tations required by DTR under any determinstic heuristic, compared to an optimal checkpointing algorithm.
Proof. We will prove this theorem by designing adversarially generated graph that forces DTR to repeatedly rematerialize evicted tensors. Our architecture simultaneously leverages the static planner's ability reorder computations, avoid repeated computation of evicted tensors.
Since DTR a dynamic algorithm, must choose which tensor to evict at time 7 based only on the portion of the graph computed up time T. Our adversarial architecture generator builds the network one node at time, choosing the next node based on the previous choice of the DTR algorithm. The construction is follows:
1. The graph begins with tensor to, which, by the behavior of DTR, must remain in memory. Tensor to has children, t through tg.
2. After step B of the computation, one of to children must no longer be in memory. Call this evicted child The next node revealed by the adversary the child of tx, causing DTR to rematerialize t..
3. The adversary continues to repeat this construction. Since to has B children, but there are only B units memory to allocate among its descendants, there must be some path from that contains resident tensors. The adversary reveals the next resident tensor the end of that path, causing DTR to rematcrialize the entire path. This repeats until we have revealed all N nodes of the graph.
An example construction of the adversarial architecture is given in Figure
Next, we analyze the computation of DTR on this graph. do this, we sum the cost computing each tensor i through f.- Consider the architecture of the final revealed network, and let I,, denote the length the path starting from E,, where j (1,..., B) so that t, direct child of fo- Since our adversary places the next node such that the entire path must be rematerialized, the total cost of computing this graph dynamically
B L
C
B
L
(Tg
+1)
2
1=1
B
Σ
8
1=1
20

---

Published as a conference paper at ICLR 2021
where hides constant factors. This sum is minimized when the L, are all equal, which gives Lj = (N - -1)BB. The cost computing all the tensors is therefore at least
C
= /B
To finish the proof, we upper bound the cost of the optimal static algorithm on this adversarial graph by exhibiting one static checkpointing algorithm and analyzing its behavior. The static algorithm may observe the entire structure the N nodes, and rearrange the computation any equivalent order.
Consider the static algorithm that computes the entire graph one path at time. That is, the algorithm first computes 1 and all its children (requiring only two units of memory, with no rematerializations). then computes t2 and all its children (again, reusing the same two units of memory), until all B paths are computed. The total cost is therefore (NN)
We that DTR requires / /B) computations compute the tensors in this graph, whereas static checkpointing algorithm would only require (NN computations. We conclude that when DTR with deterministic heuristic, there exists architecture on which it requires least ((N/B) times the runtime of a statically checkpointed evaluation.
C SIMULATOR SPECIFICATION
In this section, we provide a detailed technical specification of the DTR simulator. This includes fundamental abstractions, formal definitions of heuristics, pseudocode, runtime optimizations, and details about the log-replgying mechanism.
C.I FUNDAMENTAL ABSTRACTIONS
We designed the simulator to support computations logged from PyTTrrhh (see Sec. C.6). In PyTorch, tensor is view (containing metadata) of buffer; multiple tensors can point to single buffer. This allows us to model the various aliasing relations between tensors in Py PyToch (Paszke et al. 2017): other DL frameworks likely also use a similar representation.
Storage. At its core, DTR is runtime system for reducing memory usage. As such, storages (i.e., buffers of memory) are the underlying unit which DTR operates on. They support the following operations:
size
Storage + N: the size of the storage in bytes;
root : Storage - Tensor: the tensor whose parent operation computes the contents of the
storage (there exactly for each storage);
tensors Storage + List Linsor):: all tensors which view the storage;
resident Storage bool: true iff the storage is in memory;
locks Storage + N: the number of locks on the storage held interally by DTR (indicating
the storage is needed for pending rematerializations):
refs Storage + N: the number of external references to the storage, i.e., those held by user code.
We say a storage S is evictable if and only if resident(S). A locks(S) = 0.
Tensor: Each tensor has an associated "parent" operation op(t) which computes it (potentially along with storage(t), its underlying storage).
Each tensor also has an external reference count refs(t); in particular, each storage S has refs(S) E E(ctensors()) refs(t). The external reference count is used track whether tensor is still live in the source program whether it should be treated having been deallocated by the source
21

---

Published as a conference paper at ICLR 2021
program. Additionally, t is an alias iff + root (sotsgtorgeett), meaning that t is view of storage created by a different parent operator. For convenience, we define size t) to be 0 is an alias and size sise(storage(t) otherwise (since the metadata will likely be on CPU).
Unlike storages, a tensor is resident when storage(t) is resident and op(t) has been performed after storage(t) last became resident. This condition is denoted as defined(t). and models the behavior of our PyTorch prototype implementation where the whole tensor object is destroyed upon storage eviction (including metadata about the view, like striding and offset Thus, before operation depending on can be executed, define((!) must be satisfied, given our assumption that views of a storage must be evicted once the underlying storage has been evicted. Note that for a non-alias tensor resident(sforage(( (storage(t) if and only if (t).ned(t).
t, we have resident
Operator. An operator represents fundamental unit of computation in DTR. Operators are assumed to be pure functions of their arguments, not depending on any other external state (see Sec. C.6 for our handling of mutation). As such, each operator has an associated compute cost cost (T) E N. We assume each f has type [insor) + List|Teesen Tensor] and define imputs (f) and outputs (f) to be the input and output tensors of respectively.
C.2 FORMAL METADATA DEFINITIONS
While our abstract description of DTR in Figure Dis over tensors, the simulator operates over storages rather than tensors. Thus we must define the metadata our heuristics use over storages, providing notions of cost, staleness, and data dependencies for storages rather than for tensors.
Cost. For a given storage S, we define the compute cost of S as
cost(S)
Σ (Etensora(S)
cost(op(t)).
This is worst-case estimation: it represents the compute cost which is incurred when every tensor view of S needs be rematerialized. An alternative definition is simply cost op (roo)(S))), which may be acceptable as aliasing operations are typically much cheaper than non-aliasing.
Staleness. We estimate the staleness of S by tracking the last access time of each t tensors (S).
The last access time last_accass (tt__cceesstt is defined as the most recent time when was referenced by queued operation. Naturally, define last_ (ost_access(S)) Staleness, given the current time T, is then defined as stale (S) == T last_ (sce_accesss)
maXigtensore(S) last access
Data dependencies. The dependencies of S are the set of storages
deps(S)= = (storage(u) 3t. e tensors(S)
E inputs(op(t))) \ {S}.
Note that we exclude S since it is not true dependency (each alias tensor in tensors technically "depends" on S). Another possible approximation of the above is to simply take the dependencies of (S);; although this ignores potential dependencies of aliasing operations, it is precise if all aliasing operations depend only on S.
We now define the dependents of S the set deps ((S) consisting of all T with S deps(T). With this definition, DTR can operate over the dependency graph (V, E) where is the set of storages and (S, T) E iff S deps(T). Note that (V, E) is implicitly indexed by time T, with V being the set of non- banishe but at-least-once computed storages at T and E being the dependency relations at T.
Evicted neighborhood. The evicted neighborhood e*, as defined in Section works without modification over the storage dependency graph. We define here for completeness. Let deps, (S)
"The storage field in PyTorch tensor is immutable; in principle, we could have changed this to permit reassigning views of evicted storages to point to null and ensure the storages are rematerialized when needed, but this would have required much more extensive modifications to the codehase, which may rely on the invariant of immutable storage pointers.
22

---

Published as a conference paper at ICLR 2021
be the evicted subset of (Sps(S), and likewise for deps, (S). Now, let De and De be the transitive closures of the relations
((T [) T (S)) and ((S.T) TE deps. (S)}.
respectively. Then, ex (S) := {T' (T,S) U (s,T) D. Intuitively, e (S) is the set of evicted storages that must be resident to compute all E tensors(S), together with the set of evicted storages T that need S be resident before all E (ensors(T)) can be computed.
Relaxed (Union-Findl evicted neighborhood. Actually tracking e (S) can be computationally expensive due to the directed and changing nature of the graph. For each S, e*(S) depends on its specific ancestors and descendants, which can vary as tensors are evicted and rematerialized. An exact solution would likely involve dynamic graph connectivity data structure, which would greatly increase the complexity of the simulator's implementation.
We find an approximate solution by relaxing the definition of the evicted neighborhood. At a high level, our solution works as follows: given a storage dependency graph G = (V, E), we first forget edge directions to obtain the undirected dependency graph G. Now, let G. be the subgraph obtained by removing all resident storages (and any edges including them). Each connected component of Ge is then an evicted component, with each evicted E V belonging to exactly one component e*(T).
Importantly, track these evicted components using Union-Find (UF) data structure, which efficiently supports merging and obtaining static set metadata. Each component tracks the sum the compute costs of its elements (with the union of two components having the sum of each constituent cost). We denote the associated UF set for a storage T by T T'set, which is mutable state.
We now define the relaxed evicted neighborhood for resident storage S
E*(S) =
U T.set
T.set
TEdeps_(S)
(Ed
Note that practice, no UF unions are performed when querying this approximation. Instead, we collect and merge the set metadata separately, otherwise would erroneously merge evicted components during heuristic evaluation. This approximation reduces the worse-case time complexity of querying compute costs over the neighborhood to be linear in the number of adjacent storages, as opposed to all ancestor and descendant storages.
However, rematerializing tensor in evicted component creates a split in the component and splitting is not a supported operation on UF data structures Approaches to splitting would also need to recover the original compute costs of each set, which may require traversing the whole set if done naively. To handle splitting more efficiently, we use the following approximation: when a (previously) evicted storage S is rematerialized, we first set S.sel.cost == S.set..ost cos((S), and then assign S.set := @ (i.e., assign S to new empty UF set). Note that when storage is first computed, its evicted component is also initialized to be empty. While resident storages thus never count towards the compute cost of component, "phantom connections" between evicted storages may accumulate over time (likely depending on the connectedness of the underlying dependency graph). Despite this limitation, this approximation worked well in practice, as seen in the simulated and prototype results.
C.3 FORMAL HEURISTIC DEFINITIONS
Having defined the metadata above, (Recall that hDTR heuristics compute a score using measures of size, computational cost, and staleness and evict the tensor with the smallest score, corresponding to the intuition that the tensor evicted should be large, unlikely to be rematerialized, and cheap to rematerialize if it does need to be rematerialized.)
formally define the hpTR variants used in Sec. 4
ADTR := cost(S) + Exe (T)
sice(S) stalet(S)
This can be seen variant of the Union-Find-Shp problem, which typically requires the use of more complex data structures such link-eut trees.
23

---

Published as a conference paper at ICLR 2021
cost(S) + E Eree-(s) cost(T)
cost(S) cost*(S)
DDR (S)
size(S) staley(S)
size(S) staley (S)
Note that the simulator implementation uses the splitting approximation described above, with é (S) depending on the specific sequence of evictions and rematerializzaions. cost" (S) in the second expression is used to denote this statefulness.
(Os)(S)
DTR (S)
size(S) staley
C.4 IMPLEMENTATION DETAILS
Runtime state. In what follows, we denote the collective runtime state of the DTR simulator as R. and use the dot notation to indicate stateful reads and writes of runtime values. The simulator tracks the following runtime state:
h.euristic : (Storage, Metadata) - R, the eviction heuristic, interpreted as a score (the lowest-scored storage evicted);
R. budget N, the memory budget in bytes;
R.memory: N, the current memory usage in bytes;
R.T: N, the current clock time some unit of granularity, such as nanoseconds;
R.pool List(Storage), list of all currently evictable storages.
Eviction and banishing. To evict given storage S, we set all tensors in S' to be undefined, remove from the pool, and decrease R.memory by (Sz. Cached metadata also updated as necessary.
Banishing (permanent eviction) is slightly more subtle: in particular, it can only be done for S when d T (S) = Q. Banishing then proceeds by evicting S as above, but with the additional effect of removing S entirely from the dependency graph. Each T deps" T (S) is then locked (and effectively becomes an non-rematerializzale constant). Storages locked in this way are said to pinned (and have a special flag in the simulator), to distinguish them from those locked during rematerialization, and we permit them to banished in the future. Note that banishing can be performed on evicted when the above condition is met, in which case the eviction is skipped.
miteriaialitinnon When a tensor t is to be (re mattrriiilllized, its parents' storages are first locked by incrementing the lock count (so that they don't get evicted while they still needed) and undefined parents are recursively rematerialized. We then increment R.memory by E Euvossyyasaopthh size (performing evictions necessary), and move R.T forward by cost(op(tt) Multi-output operations must be handled carefully so as to not leak memory: we make sure to decrease R.memory by size(u') for each u' E outputs( (op(t)) that immediate freeing of doubly-computed ephemeral tensors in the Py PyTorch implementationn Lastly, locks parent storages freed and unlocked storages (including any newly rematerialized ones) are added back into R.pool.
(ttput
defined prior to the rematerialization. This models the
The simulator models non-rematerializaale constants like weights and inputs by creating dummy "constant" tensors using nullary operators with 0 cost and pinning the resulting storage. This allows the simulator to have full picture of the computation graph. Furthermore, log accurate banishing requires knowledge of constants (as PyTorch reference-count constants).
Constants.
C.5 ADDITIONAL RUNTIME OPTIMIZATIONS
Banishing and eager eviction. When the final external reference to storage is lost, we know that the underlying DL framework would have reclaimed the memory used by S. To utilize this information as opposed to doing nothing, DTR can either banish or simply evict S normally. When banishing, the runtime must first check that S has evicted dependents; if it does, then retry
24

---

Published as a conference paper at ICLR 2021
banishing each time dependent is rematerialized. Banishing has the ability to free constants, but the downside of pinning potentially exploding amounts of memory. The alternative (eager eviction) is easier to implement and simply involves evicting S normally (if possible). This prevents the problem of over-pinning memory, but with the downside that constants can never be evicted. In practice, eager evictions have allowed support lower budgets by pinning fewer values (see Sec. D.2 for details).
Caching metadata. To avoid costly recomputations of metadata during heuristic evaluations, we cache the local cost (S) for each S, only changes when new aliases are made. Additionally, for the bDTR heuristic, we avoid recomputing e" (S) at each evaluation by caching and only recom- puting after evictions or rematerializations that directly affect e* (S). Such recomputations are further optimized by tracking the evicted ancestors and descendants separately (allowing them to be recomputed independently. depending on the position of the affected storage).
C.6
MECHANISM
Log format. We logged PyTorch operations as a sequence of abstract instructions corresponding to the semantics of the actions we were easily able to instrument in the framework. Every PyTorch tensor is given unique identifier string upon creation, which is recorded and used in the log. In this section, each PyTorch tensor t corresponds to a simulator tensor [t].
The log contains the following instructions:
MEMORY(L, size): logs that é uses size memory; treated as if [t] is an alias.
ALIAS(to; f.): logs that [to]] is alias of [t,], Le., two different views of the same storage. t can either be a tensor identifier or I;; if f2 = L, then fo does not alias another tensor (t,, 's parent operation created its storage).
CALL (ipputs, outputs, cost, op): logs the operator call outputs = op(inputs, with com- pute cost cost. This instruction is followed by outputs] MEMORY and ALIAS instructions to log information about each output. Each CALL corresponds to simulator operator [[o]] with inputs [[+] i E inputs) and new simulator tensor outputs {[o] | o E outputs).
MUTATE( inputs, inputs', cost, op): logs the in-place (mutating) operator call op( (inputs) with compute cost cost, which modifies inputs' C inputs.
CONSTANT(t): logs that [t] a constant, and followed by MEMORY instruction.
COPY(L., 4): logs a new identifier to with [to] [4:]- This increments re.fs[[:.).. This happens when Python code like "x y" is called where is a Py Porch tensor and x is a fresh variable; this action neither creates a new storage nor new view but only has x point to the same view as y.
COPYFROM(to. ty): logs the PyTorch code f,- = t, where each side is an existing tensor. This decrements
increments refs([t.]). and updates [[o]] + [t.]] Intuitively, this corresponds to Python code like "x y" where y is a PyTorch tensor and was already assigned to a Py Torch tensor; in PyTorch, is mutated to match
RELEASE(t): logs the destructor of the PoTchch tensor This decrements refs([[]).
Supporting mutation. To support mutation from in-place operators, the simulator adds a "reference layer" that mutates cloned tensors, allowing for a uniform interface for all operators. Given a mutation instruction MUTATE(inpute inputs', cost, op), let inew be new unique identifier for each i E inputs', and let imputsnew = (inew E imputs"). We then proceed by treating op as a pure operator from imputs to imputs" new where each newly created simulated tensor [nee is non- aliasing and has size size (nteesarrreeeee ([=])) Lastly, we decrement refs([ and update the mapping [ + I. Intuitively, we are modeling the transformation
op(t) Tensor t' copu(t); (p
Note that in our prototype implementation, mutation of i may produce incorrect results when [i] is an alias, since the mutation layer would create a clone but aliases would still point to the old storage. Potential solutions in real implementations would be to propagate the above rewrite to all aliases of a storage (costly) or to mutate storage pointers (which would have increased the complexity of our modfications to PoTorc)).
25

---

Published as a conference paper at ICLR 2021
Output condition. All live tensors the end of log (i.e. all with refs(t) > 0) are treated as necessary outputs (namely, gradients, the loss value, and the prediction). They are thus rematerialized (if evicted) and locked to ensure they persist. This prevents the simulator from incorrectly reporting better results by evicting computed weight gradients and never rematerializing them. This permits the user to perform the weight update step outside of DTR immediately after the backward pass ends. Based on our observations of PyTorch's opt imiiieee zer gradient updates, we could also support performing these updates within DTR, since a parameter update simply performs in-place mutating additions (add_) of scaled gradients to the parameters.
D ABLATION STUDY
In this section, we present an ablation study comparing the impacts different sources information for the the lDRR heuristic. In addition to comparing the overhead in terms of additional tensor computations, we also consider the runtime overhead of different /DTR configurations in terms of the number of tensor accesses by heuristic computations and metadata updates. We also compare different eviction policies for the hDTR heuristics: ignoring deallocations, eager eviction, and banishing. These trials were performed using the same logs as in Sec. 4
D.1 DATA SOURCES
First, we will analyze the three sources of information (metadata) for the /YTR heuristic. Let us consider a parameterized version of hDTR defined ( m, (t) c(t)/[m(t)- s(t)], where a measure staleness, m is a measure of size, and is a measure of compute cost. For this study, we take and m to be the staleness and size functions defined in Appendix [ For compute cost c, compare the following alternatives (see Appendix for definitions): the full e", the approximation ě* and the local cost (cost of the parent operator only). We allow each measure to be entirely ablated no).
(e.g., s(t) = 1, which we denote
In the following figures, we specifically have s, m E {yes, no) and E {e*, EqClass, local, no). Each figure fixes choice of c, varying and m.
2.00
Inception (64)
Transformer (10)
U-Nett (o)
TrecLSTM
1.50
125
1.00
2.00
Compute
1.75
2212224
Dens
rour divinoee D13 Sczuencelapph
12 caps, 2120512
1550
125
1.00
3.7
0.9
0.3
0.5
0.3
Memory
stoleness
.iize, to stalenea
no atalenass
stalemmes
Figure Results for fixed e", varying and m.
The general trend shown in Figures 7 8 9 10 that higher metadata complexity (corresponding to more precise notions of the evicted noighborhood) enables more savings, while staleness and size are required for acceptable computational overhead. staleness and size depends on the specific model architecture. For example, cost and size alone each do far better than using both cost and staleness for the static models (DenseNet, ResNet, UNet), whereas the opposite is true for the dynamic models. This may be due to model depth or the distribution of tensor sizes or to the increasing impact of individual checkpoints at lower budgets; further research may shed more light the influence of model-specific characteristics like these. Additionally, may note that the e* approximate cost performs comparably to the e" exact cost while requiring less information, validating our claim that the equivalence classes are useful approximation.
interesting to note that the importance of
26

---

Published as a conference paper at ICLR 2021
2,00
InceptionVe
Transformer (10)
U-Net
TreeLSTH
1.75
4106006
1.50
Compute Overhead
1.00
1.75
2246224
2446204
Tou disauue 112 besiexelee
13 supo 5320512
1.0
125
1.00
0.3
0.7
Eq Casas pioe,
EqCass, $26,
staleness
staleness
Figure 8: Results for fixed c EqClass, varying and m.
2.00
InceptionV& (64)
Transformer (10)
U-Net (6)
TreeLSTM
1.50
Computs Overhead
1.00
1.75
2242224
your divinkee 113
12 mapa, 2120512
Sezuenoolangh
1.50
125
1.00
0.5
0.3
0.7
Memory
Figure 9: Results for fixed local, varying s and m.
2.00
IncoptionVV (64)
Transformer (10)
U-Nett (6)
TreeLSTM
1.50
Compute Overhess
1.00
1.75
2240224
2246224
siip $123512
Sequace 12
230
1.5
1.00
0.7
0.3
3.7
0.9
dec, staleness
cost
staleness
cocl,
leness
s20. ataleness (randon))
Figure 10: Results for fixed c no, varying and m.
In general, the best-performing of these heuristics were those with non-ablated choices of s, m, and c, hence our choosing the figg variants with e", ē*, and local cost (/DTTR, ADTR. and A ADDR. respectively) for the evaluation in Sec. 4
D.2 BANISHING AND DEALLOCATIONS
For the following trial, we compared the /DTR heuristic with banishing (permanent removal) against that with eager evictions, described in Appendix C.5 We also compare both deallocation-awan approaches against simply ignoring deallocations. We only used e* cost because it performed much better than local cost and because it would have been more complicated to update the definition of 2* to account for banished neighbors.
The results shown in Figure 11
27

---

Published as a conference paper at ICLR 2021
2.00
InceptionVE
Transformer (10)
U-Net
TreeLSTH
1.75
4106008
1.50
Compute Overhead
1,00
1.75
2246224
2446204
Tox dinauue 112
13 supo 5220812
150
125
1.00
0.9
0.3
Banlshing
Eager exidtion
Figure 11: Results for the hDTR heuristic, comparing banishing and eager evictions.
As the curves show, banishing is not able to achieve the same budgets across most models tested as eager eviction. For UNet, the difference is large: banishing can only support 90% of the baseline budget (and OOMs at 0.8 ratio), while eager eviction can support 50% of the baseline budget. However, banishing still attains low budgets on most models, even obtaining better computational overhead under the same budget and savings for ResNet. Since banishing potentially allows for greatly lowered runtime overhead, implementations of DTR can consider conditionally enabling it in situations where the tradeoff is more desirable.
Compared to ignoring deallocations, both banishing and eager eviction obtain noticeably lower rematerialization overhead. This shows that valuable information is captured by deallocations, and that DTR can make good use of it.
D.3 RUNTIME OVERHEAD
For this experiment, we tracked the number storage (see Appendix C.I accesses made during evaluations of heuristics and maintenance of metadata. We chose this metric over wall-clock time, since our Python implementation of the simulator is not heavily optimized and may not accurately correspond to the real performance of the runtime. Storage accesses, on the other hand, do reflect operations that would be performed by a real implementation. For the İDTR heuristic, this included each storage visited during the updating and rebuilding procedures for maintaining e" for resident storages. For the hpiR heuristic, this included each storage visited whenever the Union-Find data structure was traversed for each evicted component (which occurs mainly during merging and when reading the compute cost). The /oor heuristic does not need to maintain any non-local metadata. For all heuristics, each heuristic evaluation counted as one storage access.
Inception 10
10*
20"
201
301
0.3
ResiNet-32 32 (56)
10"
DenpeNot- 0643
10a
LSTM (512)
Uérelled GAN
107
20*
Sequsnce angos
204
10a
0.
3.
0.9
0.3
0.7
0.5
3.7
0.9
0.1
05
0.7
0.9
Memory
16t
11367
Figure 12: Total storages accesses incurred by heuristic evaluations and metadata maintenance, compared across different memory ratios, for the 3 main MoTR variants.
As Figure 12 shows, the accesses made by each heuristic are generally separated by at least an order of magnitude. This confirms our intuitions about the runtime overhead of each heuristic, and supports
28

---

Published as a conference paper at ICLR 2021
our choice of RDTR good middle ground (in terms of both runtime and computational overhead). However, these overhead figures could be improved with better-optimized implementations of the heuristics, our implementation recomputes heuristies often, even when it may be possible to store the scores for tensors and maintain them in a sorted order. (Reformulating staleness to avoid having to use the current time might help.) Using persistent data structures that can be incrementally updated and maintain a sorted order will make these heuristics much more efficient, though this would also increase the complexity of the implementation.
E PROTOTYPE IMPLEMENTATION
E. INTEGRATION INTO PYTORCII
To avoid modifying PyTorch's core systems, our DTR prototype is implemented as a wrapper over PyTorch' existing tensor implementations. Namely, we add new tensor representation into PyTorch called a Checkpoi nttpoinnTio Tensor, which is simply a wrapper over an existing PyTorch tensor that additionally tracks the tensor's parent operation and other metadata (such as the last access time and the cost of the parent operation, which is timed when the tensor is first created) and registers the tensor in the DTR runtime system. Timing operators for metadata purposes simply uses the system clock, hence to guarantee the correctness of these operator times, we force Py Torch into synchronous execution mode (which ensures that GPU operators are performed synchronously); we found that DTR was still able to execute models on greatly reduced memory budgets without turning on synchronous execution mode, even though this should skew DTR's 's recorded operator times.
For evictions, Checkpoint Tensors are capable of freeing their underlying tensor representa- tion from memory; they keep a closure for replaying the parent operation, which the runtime can invoke when the tensor must be rematerialized. To handle deallocations by the original program, Tensors also report increments and decrements to the reference count of the un- derlying tensor to the DTR runtime. We add method to tensors called "checkpoint ()" " that lifts any tensor into a SeteckkoinnTenssee underlying tensor from Checkpoing Tensor, Tenssor rematerializing it if necessary (we use the latter in our trials to ensure the loss and output are in memory at the end).
and a method "decheckpoim () that extracts the
Our modified version of PyTorch dispatches any operation involving a CheckpoirtTennsor to a specific implementation for Checkpoint itnkkoornnneessrrr. this is the same mechanism that Py- Torch uses, for example, to dispatch operations on GPU-managed tensors to CUDA imple- mentations. Specifically, whenever PyTorch encounters an operator where an argument is a Checkpoint its dispatch mechanism searches for a specific overload of that operator Since a Checkpoint Tensor simply wraps the underlying Py-
for Checkpoint Torch tensor, adding CheckpeintTensor implementations for operators simply requires invok- ing the operator's s existing implementation for the underlying tensor and wrapping the result in Checkpoi i nt Tensor. These overloads were essentially boilerplate code and it is likely possible to generate them automatically. As far as PyTorch's 's dispatch system is concerned, all tensor accesses occur through operators, so updating metadata like access time only reqires invoking the DTR rumtime inside the Checkpoi i nt Tensor operator overloads.
The DTR runtime is simply a singleton that keeps a pool of all Checkpo nt Tensors created since the start of the program. The runtime is also responsible for maintaining the equivalence class data structure needed for A described in Appendix C.I (updated each time a Checkpoint Tensor Tensor is evicted or rematerialized). Before each Checkpoint Tensor operation, the DTR runtime checks whether the memory budget has been exceeded; if it has, the runtime searches over the pool of www.eoooonnntntooooeo
computing the heuristic score (h for each using their metadata, and evicting the least- scoring until either it is not possible to evict any more tensors or the budget has been met. bN. this means that the prototype permits exceeding the budget by exactly one tensor allocation. In principle, we can correct this by inserting a callback into PyTorch's 's GPU memory manager to call the DTR runtime as soon as an allocation is requested; we did not do this to simplify our implementation.) This method of searching is very simplistic; it is likely that redundant heuristic computations can be removed using data structures to keep Checkpo i Tensors in a sorted order and incrementally update metadata, but the optimizations discussed below in Appendix E.2 were very simple and helped to reduce some of the overhead from this naive method. The DTR runtime is also responsible for implementing the logging mechanism described in Appendix C.6. this is
29

---

Published as a conference paper at ICLR 2021
accomplished by simply writing JSON records of events intercepted by the runtime (operator calls, reference count increments and decrements, etc.) to file.
The DTR prototype supports PyTorch's implementation details like in in-place mutations, aliasing, and multiple operator outputs, which are all discussed in Paszke ct al. (2017), using the same methods as the DTR simulator (see Appendix [C). As in Appendix|C.66] the DTR prototype supports PyTorch operators that perform in- in-place mutations by introducing a copy-on-write mutation layer: The mutating operator is made pure (and therefore infinitely replayable) by copying the source tensor for the mutation and mutating the copy. (Similarly, impure operators like batchnorn and dr dropout are made pure by treating state like the PRNG seed as part of the imput to the operators and the updated state as part of their output.) The DTR runtime performs these copies for Checkpoinss Tensor operator overloads to mutating operators. To support operators whose results are aliases of their arguments, the DTR runtime groups together all Checkpo i Tensors whose underlying tensors are aliases of each other into alias pools. When member of an alias pool is evicted, all members of the alias pool are treated as evicted; aliases are, however, rematerialized separately, only as they are needed. For Checkpoint Tensors produced by multi-output operations, the DTR runtime allows them to be evicted separately but ensures that they are rematerialized together.
E.2 RUNTIME OPTIMIZATIONS
Searching for tensors to evict is significant source of overhead for DTR's runtime because the runtime recomputes each tensor's 's staleness and equivalence class cost upon each eviction, rather than storing and incrementally updating this information. In principle, we could reduce this portion of the overhead by using more complex data structures to maintain an ordering of the tensors to avoid searching, though this would greatly increase the complexity of our implementation. As a simpler means of reducing the DTR runtime's overhead from searching and computing heuristic scores, we added two approximate optimizations to reduce the search space: ignoring small tensors (less than 1 1% of the average size) and only searching over random sample of /n tensors from the pool of evictable tensors. This greatly reduces the number of tensors that the runtime needs to check upon evicting. Even though this improves the search overhead considerably, searching and computing costs still present considerable DTR-specific overhead, as the profiling breakdown in Figure 4 shows. Additionally, random sampling caused occasional failures low budgets or very large inputs due to excluding good eviction candidates from the search space, which led us to deactivate that optimization in certain trials. (At low budgets, individual eviction choices are very impactful, so removing tensors from the search space completely at random can dramatically affect the results.)
There are also several possible sources of runtime overhead that could potentially be improved by making deeper modifications to PyTorch's 'ss core systems. For example, we introduced an overload layer that results in many more layers of callbacks. The mutation layer also clones tensors (even though it frees the necessary space immediately), resulting in additional overhead. Further modifications to the framework could allow for more optimizations, particularly by reducing the number of heap allocations and conversions between tuples and lists. PyTorch's 's define-by-rut nature and shallow embedding into Python also meant that much of DTR's 's metadata, such as the parent operator of a tensor, needed to be computed at run time (such by creating closure). In other frameworks that feature a compilation step, such as Glow (Rotem et al.) 2018), it may be possible to eliminate much of this overhead by generating these structures in compiler pass. We may also note that all the bookkeeping for DTTR takes place on CPU while operators are generally offloaded to other devices, so an implementation could interleave these updates with GPU operations.
E.3 HANDLING ERRORS IN TRIALS
As discussed in Table and Figure the DTR prototype encountered errors on certain models when running on low budgets or on large input sizes. These errors were primarily CUDA out-of-memory errors (OOMs), but in some cases, the trial simply hung, neither crashing nor terminating. For CUDA OOMs, disabling the random sampling optimization described in Appendix E.2 eliminated the errors in most cases, suggesting that the OOMs were due to excluding useful eviction candidates. For the hanging trials, trapped in a very deep recursive rematerialization, as occurred in some of the simulated trials on certain heuristics) an infinite loop or deadlock elsewhere in PyTorch; we can investigate the cause
not able to determine whether the root cause was DTR thrashing (being
30

---

Published as a conference paper at ICLR 2021
by further instrumenting the implementation, but have been unable to consistently reproduce hanging trials and they seem to occur less frequently than OOMs.
In the largest two batch sizes for UNet in Table [ disabling sampling did not climinate all OOMs hanging trials. Thus, for the large-input trials in Table we employed a procedure for retrying upon encountering an OOM untimed "warm-up" trials to allow for CUDA initialization and caches to be populated and then begin timing the trials. If a trial raises CUDA OOM or hangs (which we define as taking twice as long as the trial before it), we keep the measured times from that point in the trial and then restart (doing another warm warm-up), up)). collecting the remaining number of measurements. Restarting the measurement run was the only way to ensure that all memory allocated during the trial would be collected in the event of an OOM (attempts to proceed simply by resetting the PyTorch allocator's cache resulted in memory accumulating between trials regardless). Our experimental setup automates this process of retrying failed trials and reports the total number of retries. Note that we treat failures during warm-up runs the same as failures in timed runs, since recovering from an OOM would require exiting the process running PyTorch and reinitializing CUDA. the Table results, there was failed run for UNet on batch size and 10 failures on batch size 10; most of the latter were during warm-up runs.
hang. First (as with all other GPU measurements), we perform some
possible reason for the occasional failed trails UNet may be variance in operator timings, which affect the metadata and may be influencing rematerialization decisions. One way to control for this possibility in a static model like UNet would be to use a DTR simulation to produce a static rematerialization schedule and therefore have known, safe execution schedule for operators. For dynamic model, a static plan is not an option, but variations in operator timings could be reduced by using fixed cost model for operators instead of timing them dynamically. That is, the DTR heuristics employed could be defined to use proxy measures that are less subject to variation (e.g., defining staleness in terms of counter incremented by operations rather than wall-clock time) less likely to be influenced by specific system implementation details in order to have more predictable and reproducible behavior.
31


Reviews:
[{"review_id": "h0fmv_6rYmz", "reviewer": "AnonReviewer3", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Claims:\n\n-Better eviction heuristics provide noticeable improvement over earlier dynamic rematerialization schemes - up to 50-80% memory savings at the cost of up to 75% increase in computation\n\n-Seamless implementation with PyTorch; user does not need to change their code - highly impactful if incorporated into PyTorch\n\n-Upper bound analysis shows only O(N) operations are needed (constant factor more than without rematerialization)\n\n-Several eviction heuristics are studied, including from previous literature. Mathematical formalism unifies these heuristics and clearly describes their relationships\n\nPros:\n\n-It seems that previous rematerialization schemes provide up to 50% memory savings (i.e. can run 2x larger models), but the proposed work goes up to 80% memory savings (i.e. can run 5x larger models). If this is true, then the ability to run 5x larger models on the same hardware shifts rematerialization from a relatively minor optimization to a game-changing feature\n\n-Running time overhead below 100%, which is an acceptable tradeoff (authors consider runtime overhead >100% to be thrashing, which is sensible)\n\n-More so than typical papers, the design motivations are factually and intuitively explained, and performance claims are not exaggerated but rather put into perspective. The writing strikes me as convincing.\n\n-Comprehensive experiment design with 8 models and 7 eviction heuristics, covering RNNs, CNNs, Residual networks, Transformers, and unconventional models such as Unrolled GAN. This leaves little doubt as to the proposed method's effectiveness\n\nCons:\n\n-The paper did not study rematerialization in at least the multi-GPU, if not distributed-parallel setting. I think the latter can be excused given the typical scope of an ICLR paper, but I was expecting the former and was surprised to find the experiment design only involved one GPU. At minimum, a multi-GPU experiment is needed to confirm the implementation works with multi-GPU setups, if not distributed-parallel setups.\n\nQuestions:\n\n-The paper does not talk about datasets used. This is appropriate if all models in the experiment suite do not have dynamic structure, since dynamic structures could (depending on the implementation) cause the computational graph to differ with different datasets - e.g. I imagine this could happen with NLP and BiLSTM or recursive neural networks. Can the authors clarify if they limited their models to those with static (fixed computational graph) structures? How would the proposed method behave on models with dynamic structures that change with the input data?\n\nSuggestions for improvement:\n\n-The authors call out the relationship between rematerialization and tensor swapping only in the related work. While this is better than not bringing it up at all, it can be argued that since rematerialization and swapping are both dynamic memory management techniques, swapping should at least be mentioned in the introduction. Even if swapping is out of scope of the paper, I think some discussion on how DTR and swapping could be combined would make the paper even stronger.\n\n-Although the paper does explain and provide experiments as to why dynamic rematerialization is superior to static rematerialization, I think the argument could be made much more convincing if the authors showed evidence that static rematerialization does not work on certain dynamic models, limiting its applicability. For example, a statement to the effect of e.g. \"Checkmate doesn't work on model X, and here is why...\" could be instructive.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "aXFcfjqDjqp", "reviewer": "AnonReviewer4", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nMerits of the paper:\n- Address an important practical problem on how and when to perform checkpointing during DL training.\n- Cover a pretty comprehensive study across theory, simulation and system implementation.\n- The suggested system implementation looks simple and clean.\n- Clearly written paper, which is easy to understand and follow.\n\nPlaces to improve:\n- The need of having dynamic approach in this area is not very well motivated.  Since the computation in most of DL models is repetitive over iterations, static approach would work pretty well.  Furthermore, since many models take long to train, spending minutes on analyzing static graph and obtaining an optimal solution seems to be time well spent, comparing with a suboptimal dynamic solution.  The motivation of developing dynamic approach needs to be strengthened.  \n\n- Although some comparison with related work is done by simulation, really system evaluation is relatively weak - it only compares with the strawman PyTorch models without checkpointing.   It would be a lot more convincing with the results comparing with basic checkpointing approach (e.g., layer wise checkpointing) and some related work.\n\n- For some of the reported models, like Unet, the approach reduces memory footprint and helps increase batch size, which however reduces  throughput.  These are probably not good examples to show the benefits of the approach.  But it might be beneficial to point out the underlying reasons for decreased throughput so readers know when the approach would/wouldn't work well and why. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "Vyzy7fCTQe", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: This paper proposed a simple yet effective greedy algorithm with a new heuristics on checkpointing deep learning models so that people could train large model with restricted GPU memory budgets. The proposed method operates in an online setting and do not need static analysis of computation graph, thus could be used for both static and dynamic models. In a restricted model setting of linear forward network and equal space and time cost for each node, the author proves the proposed method could reach the same bound on tensor operation and memory budget with previous static checkpointing methods. The author also establish a theorem on tensor operation numbers between the proposed dynamical method and an optimal static checkpointing algorithm. In experiment, the author compared the proposed method with static techniques including the optimal Checkmate tool of Jain et al. (2020), showing the proposed method gives competitive performance without static model analysis in prior. The author also compared the proposed heuristics with prior arts on several static and dynamic models. Finally, the author described a prototype of PyTorch implementation of the proposed method. \n\nPros: \n1. While goes under a limited setting, the theoretic analysis on the tensor operation and memory budget bound of the proposed method, as well as on the relationship between the proposed method and optimal static analysis method is novel and interesting. The experiment also shows the competitiveness of the proposed method by comparing to static methods. \n2. The author does a great job explaining the idea, concepts, procedures and experiments. \n\nCons: \n1. The author provides the comparison between the proposed heuristics and others with the same greedy algorithm, but it seems not to have the full comparison to other dynamic checkpointing approach(e.g. Peng et al. (2020) ). Although experiments in the paper shows competitive results with static model, the main use cases of the proposed method might still be in dynamic models as in normal static model use cases, the time overhead of static analysis could be ignored compared to actual model training time.\n2. As the proposed heuristics bears some similarity with the one used in Peng et al. (2020), it would be more convincing to also have an ablation study of replacing the heuristics used in Peng et al. (2020) with the proposed one. \n\nReferences:\n\n[1] Jain, Paras, et al. \"Checkmate: Breaking the memory wall with optimal tensor rematerialization.\" Proceedings of Machine Learning and Systems 2 (2020): 497-511.\n\n[2] Peng, Xuan, et al. \"Capuchin: Tensor-based GPU Memory Management for Deep Learning.\" Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems. 2020.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "TId0UKb-YCS", "reviewer": "AnonReviewer2", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nPrior such work all required static analysis and planning of the network - and hence were of limited use. Significant contribution of this work is summed up at the end of Sec 4.3. It achieves 'as good' results as prior state-of-the-art (Checkmate, published at MLSys) - but without any prior knowledge of the model. This significantly increases the practical significance of this work. NeuIPS 2019 work of Kumar is the other often-cited work, also based on static planning, where the authors had already noted the primary limitation keeping it from getting adopted: \"algorithm yields asymptotically better schedules, the schedule length and memory depend exponentially on the path width\". \n\nI am also happy to see authors offer hardware detail of their experimental platform (Figs 2 and 4).  And PyTorch software prototype should make it easier to follow through in other frameworks. There is a decent variety in the chosen set of benchmarks as well.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents an online algorithm for dynamic tensor rematerialization.  The theoretic analysis on the tensor operation and memory budget bound of the proposed method, as well as on the relationship between the proposed method and optimal static analysis method is novel and interesting.  It covers a pretty comprehensive study across theory, simulation and system implementation. In addition, the paper is well written. 
Decision: Accept (Spotlight)