PaperlayerView exact revision on arXiv

Source-linked AI summary

Dynamic Tensor Rematerialization

Marisa Kirisame, Steven Lyubomirsky, Altan Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, Zachary Tatlock

arXiv:2006.09616v4cs.LGcs.PLstat.ML

TL;DR

Current checkpointing methods plan recomputations offline, limiting their use with dynamic models and requiring advance model analysis. DTR makes checkpointing an online runtime decision; simulations show performance close to optimal static checkpointing, and a PyTorch prototype demonstrates modest integration changes.

  • Problem

    Existing deep-learning checkpointing techniques statically plan recomputations offline and require an initial model-analysis stage.

    Source: §1 · p. 1

  • Method

    DTR is a greedy online checkpointing runtime that uses tensor and operator metadata to choose activations for eviction and later rematerialization without advance model knowledge.

    Source: §1 · p. 1

  • Results

    DTR’s heuristics perform remarkably close to Checkmate’s optimal solutions, while hLRU outperforms the static baselines.

    Source: §4.3 · p. 6

  • Takeaways & Limitations

    A PyTorch prototype shows DTR can be incorporated with modest, non-invasive framework changes.

    Source: §1 · p. 2

    DTR cannot always match optimal static checkpointing on arbitrary networks because it cannot inspect or reorder the network.

    Source: §3 · p. 4

Abstract

from arXiv · show

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.

1 INTRODUCTION

Memory demands make training and exploration difficult on memory-limited devices, while existing checkpointing requires offline planning. DTR replaces that planning with a greedy online runtime and is evaluated through theoretical bounds, simulations, and a PyTorch prototype.

1 INTRODUCTION

  • Growing model memory demands impede training on memory-limited devices and constrain exploration of memory-intensive architectures and techniques.Checkpointing addresses this by freeing activations and recomputing them on demand, but current approaches plan recomputations offline and require model analysis.Source: §1 · p. 1
  • DTR uses runtime tensor and operator metadata to decide which activations to free and later recompute, without advance model knowledge.Its online approach supports dynamic models, including data-dependent control flow, unlike static planners that require unrolling and planning for distinct inputs.Source: §1 · p. 1
  • The paper contributes a linear-network operation bound, near-optimal simulated performance, and a PyTorch prototype requiring modest framework modifications.It also formalizes checkpointing as online rematerialization and defines a greedy algorithm parameterized by caching-inspired heuristics.Source: §1 · p. 2

2 DYNAMIC TENSOR REMATERIALIZATION

DTR intercepts tensor lifecycle events and uses eviction heuristics to manage memory, recomputing evicted tensors when needed. Its approach depends on explicit operator assumptions and balances eviction value against rematerialization cost and runtime overhead.

2 DYNAMIC TENSOR REMATERIALIZATION

  • When allocation exceeds available memory, DTR evicts resident tensors; access to an evicted tensor triggers replay of its producing operation, recursively rematerializing missing arguments.Constants cannot be evicted, and rematerialization can trigger further evictions when memory is exhausted.Source: §2 · pp. 2–3
  • DTR assumes opaque operators, operator-produced or constant tensors, single-tensor outputs, and pure deterministic operators; an operation that cannot fit may make rematerialization fail.The heuristic affects failure likelihood because eviction choices can create deeply nested rematerializations requiring many tensors to remain resident.Source: §2 · p. 3
  • The hDTR heuristic balances tensor staleness, memory reclaimed, and projected rematerialization cost, evicting the tensor with the minimum cost-to-memory-and-staleness score.Its projected cost includes costs across the evicted neighborhood, accounting for potential recursive rematerializations.Source: §2 · p. 3
  • Tracking evicted neighborhoods exactly can add runtime costs, so DTR approximates them with an undirected union-find representation that performs nearly as well with fewer metadata accesses.The approximation uses a constant-time split approximation and can introduce phantom dependencies.Source: §2 · p. 3
  • Eagerly evicting deallocated tensors reduces future runtime overhead, while permanently banishing them can save memory but prevent later rematerialization of their children.Evicting rather than banishing deallocated tensors preserves potential evictions but adds management overhead and keeps constants in memory.Source: §2 · pp. 3–4

3 FORMAL BOUNDS

On linear feedforward networks, DTR can achieve a linear tensor-operation bound under a constrained memory budget. That guarantee does not extend to arbitrary networks, where an adversarial operation order can make DTR substantially worse than optimal static checkpointing.

3 FORMAL BOUNDS

  • DTR executes a forward and backward pass on an N-node linear feedforward network in O(N) tensor operations under the theorem’s memory-budget condition, matching the static bound up to constants.The proof uses a particular eviction heuristic and assumes unit space and time cost per tensor.Source: §3 · p. 4
  • The linear-network result does not imply a general match to optimal static checkpointing because DTR cannot inspect or reorder arbitrary networks.For the linear case, evenly spaced resident checkpoints limit successive rematerialization, and backward-pass overhead between checkpoints sums to a constant.Source: §3 · p. 4
  • For any deterministic heuristic, some N-node network with budget B ≤ N requires DTR to perform Ω(N/B) times more tensor computations than optimal static checkpointing.The adversarial construction schedules operations to force rematerialization of fully evicted paths, while static planning can reorder the work.Source: §3 · p. 4

4 HEURISTIC EVALUATION

Across simulated models, DTR operates under restricted memory with trade-offs between eviction quality and runtime overhead, and its heuristics approach or exceed static checkpointing performance. The evaluation compares heuristic variants, memory-budget behavior, and static methods.

4.1 HEURISTICS EXAMINED

  • The evaluation compares evicted-neighborhood hDTR, its union-find approximation heqDTR, local-cost hlocalDTR, and LRU, size, MSPS, and random baselines.heqDTR closely matches hDTR while requiring fewer operations per eviction and rematerialization.Source: §4.1 · p. 5

4.2 COMPARING DTR ACROSS HEURISTICS

  • Across Figure 2’s model and budget comparisons, DTR trains with a small fraction of normal memory, and heuristics using more rematerialization information operate at lower budgets with fewer rematerializations.The figure plots computational slowdown against memory budgets as fractions of original peak use; complex heuristics add runtime overhead, with hDTR incurring up to two orders of magnitude more metadata accesses per batch than heqDTR.Source: §4.2 · pp. 5–6, Figure 2 · p. 6
  • Even simple hLRU typically enables training with 30% less memory, while more complex heuristics trade additional runtime overhead for lower budgets and fewer rematerializations.The authors suggest heqDTR may be more effective in practice because it closely matches hDTR with much less runtime overhead.Source: §4.2 · p. 6

4.3 COMPARING DTR TO STATIC TECHNIQUES

  • DTR’s hDTR and heqDTR approaches perform remarkably close to Checkmate’s optimal solutions, while hLRU outperforms the static baselines.Checkmate uses ahead-of-time model knowledge and seconds or minutes per budget; DTR finds comparable solutions dynamically in milliseconds without that knowledge.Source: §4.3 · p. 6

5 PROTOTYPE IMPLEMENTATION

The PyTorch prototype trains models at larger input sizes than unmodified PyTorch in several cases, including dynamic TreeLSTM, while profiling measures runtime overhead across models and memory budgets. Its implementation uses modest framework changes, and DTR’s operator overhead is competitive with Checkmate’s.

Prototype implementation

  • The DTR core uses 1,161 lines of code and avoids deep changes to PyTorch memory management or tensor abstractions.The remaining 2,647 lines of changes are primarily boilerplate operator overloads that dispatch tensor operations through DTR’s core logic.Source: §5 · p. 6

Operator-overhead comparison

  • DTR’s operator overhead is competitive with Checkmate’s ILP-produced optimal rematerialization schedule, alongside Treeverse and Chen algorithm variants.The comparison adapts Jain et al. (2020)’s evaluation artifact to translate graph representations for the DTR simulator.Source: Figure 3 · p. 7

Empirical evaluation

  • DTR trains models on larger input sizes than unmodified PyTorch in several cases, including dynamic TreeLSTM, with Table 1 reporting median batch times.The table compares supported input sizes and execution times across models.Source: §5 · p. 7, Table 1 · p. 8

Runtime profiling

  • Figure 4 profiles eight models across restricted memory budgets, and DTR’s search and heuristic-computation overhead remains low for most budgets.The prototype’s overhead breakdown separates cost compute and eviction-loop time from unprofiled time.Source: Figure 4 · p. 7, §5 · p. 8
  • The prototype supports unusual features including Unrolled GAN’s higher-order partial derivatives and extensive Python reflection, illustrating DTR’s generality.The large unprofiled-time share in Unrolled GAN is likely due to its extensive Python reflection.Source: §5 · p. 8

6 RELATED WORK

Prior checkpointing work includes static graph-based plans, segment recomputation, graph-theoretic analysis, and ILP optimization. DTR differs by planning online for arbitrary dynamic control flow, while runtime planning adds overhead; tensor swapping and reversible model designs are complementary approaches.

Checkpointing in Reverse-mode Automatic Differentiation

  • Reverse-mode AD checkpointing reduces stored tape values by recomputing segments, with Treeverse using binomial partitioning for logarithmic space and computation growth.Later methods extend this approach to arbitrary control flow by inserting checkpoints at compile time, but do not use dynamically gathered information.Source: §6 · p. 8

Checkpointing in DL

  • Static deep-learning methods recompute network segments or individual activations, with graph analyses and ILP used to plan rematerialization.Chen et al. (2016) give O(N) memory with one extra forward pass, or O(log N) memory with O(N log N) additional tensor operations; recurrent-network methods support some dynamism.Source: §6 · p. 8
  • DTR handles arbitrary dynamic control flow online, unlike static plans, but its runtime heuristics face overhead constraints that limit feasible analysis.Static greedy schemes such as Chen et al. (2016) and GreedyRemat use tensor sizes only.Source: §6 · p. 8

DL Memory Managers

  • Swapping systems and rematerialization are complementary, but combining swapping with DTR may complicate overlapping communication and computation and require lookahead.Possible designs include replacing an existing system’s rematerialization scheme or treating swapped-out values as unevictable.Source: §6 · p. 9

Memory-Efficient DL Model Designs

  • Reversible layers permit recomputing a forward value during backpropagation from the following layer’s result, and have been used in ResNet and Transformer variants.These model designs can be used alongside checkpointing and swapping approaches.Source: §6 · p. 9

7 CONCLUSION

The conclusion presents DTR as a customizable online checkpointing approach that closely matches optimal checkpointing in simulation and supports dynamic models. The appendix analyzes a reduced heuristic on a unit-cost linear network, formalizing liveness and tensor eviction in that setting.

7 CONCLUSION

  • DTR supports a broad range of applications without ahead-of-time analyses, manual annotations, or model modifications, and closely matches optimal checkpointing in simulation.The PyTorch prototype interposes on tensor allocations and operator calls while collecting lightweight tensor metadata.Source: §7 · p. 9

A PROOF OF THEOREM 3.1

  • The appendix proves O(N) runtime for DTR on a linear feed-forward network with uniform operator compute and memory costs, using a reduced heuristic.The proof addresses Theorem 3.1.Source: A PROOF OF THEOREM 3.1 · p. 12

A.1 NETWORK DEFINITION

  • The proof models N operators and their gradient operators with unit tensor size and unit compute cost, treating the input as resident and outside active-memory accounting.The input tensor is t0; each operator fi computes fi(ti−1), and each gradient operator computes from its forward input and succeeding gradient.Source: §A.1 · p. 12

A.2 LIVENESS AND BANISHING

  • A tensor is live while a pending operation needs it; once no longer live and its children are resident or banished, it can be banished and its memory reclaimed.Liveness can be analyzed statically for static models or by reference counting for dynamic models, with information delivered online through deallocation events.Source: §A.2 · pp. 12–13

A.3 HEURISTIC DEFINITION

  • The reduced heuristic omits tensor staleness, and under the simplified compute and memory constraints its score is |e*(t)| + 1.Uncomputed tensors are excluded because their existence is not yet known during dynamic execution.Source: §A.3 · p. 13

A.4 PROOF OF THEOREM 3.1

The proof establishes O(N) total tensor operations for the analyzed linear feedforward network. It bounds rematerialization by using persistent checkpoints and the additional memory freed as backward processing advances.

PHASE 1: FORWARD PASS

  • Evenly spaced forward-pass checkpoints remain resident through the backward pass until banishment, limiting consecutive rematerializations.The proof identifies these checkpoints as the basis for controlling rematerialization work.Source: §A.4 · p. 13, PHASE 1: FORWARD PASS · p. 14

PHASE 2: BACKWARD PASS

  • Backward gradient computations use a pinned gradient input and a forward tensor whose evicted path is not too long, bounding rematerialization for both inputs.The proof says the next gradient remains pinned until banishment, while the forward input’s evicted path is bounded by a later lemma.Source: PHASE 2: BACKWARD PASS · p. 14

2. The gap between neighboring tensors in C satisfies

  • Processing each checkpoint group frees one memory unit, which supports intermediate checkpoints and reduces rematerialization overhead in later groups.The proof uses the additional free memory to process gradients in later groups with less overhead.Source: §2 · p. 15

A.4 PROOF OF THEOREM 3.1

  • DTR’s forward and backward passes together require O(N) tensor operations, matching the order of the Θ(N) cost with unlimited memory.The backward-pass bound is O(N), and the proof adds the forward-pass cost to obtain the total.Source: §A.4 · p. 13, §2 · p. 15

A.5 PROOFS OF INTERMEDIATE RESULTS

The intermediate proofs bound checkpoint gaps under h_e* across the backward pass: an initial phase establishes the gap bound, and induction preserves it. They also analyze the rematerialization cost of processing gradients between checkpoints.

A.5 PROOFS OF INTERMEDIATE RESULTS

  • During the remaining N−2 backward gradient computations, h_e* keeps checkpoint gaps within 4(N−2)/(B−1), so C remains resident until tensors are banished.The proof establishes this bound after the first two gradient computations and maintains it through the rest of the pass.Source: §A.5 · p. 17
  • The first two backward gradients evict at most one tensor from S, leave two free memory units, and produce a maximum checkpoint gap no greater than 2LS.This establishes the base conditions for the later gradient computations.Source: §A.5 · p. 17
  • For subsequent gradients, induction uses two free units to rematerialize the next tensor while preserving the 2LS gap bound.After each gradient, the algorithm can banish two tensors to restore the two free units.Source: §A.5 · pp. 17–18
  • For the first gradient group, which has no spare memory for checkpoints, successive gradients rematerialize progressively shorter portions of the group, yielding a total-cost bound.The first gradient rematerializes the full group; each following gradient requires at most one fewer intermediate tensor.Source: §A.5 · p. 18
  • With k+2 free memory units and a resident starting tensor, the intermediate checkpointing argument bounds the gap after computing tensor j+i by 2(i−1)/(M−1).Partitioning the tensors into M segments yields a short adjacent pair; h_e* chooses an eviction that preserves the bound.Source: §A.5 · pp. 16, 19

2. The gap between neighboring tensors in C satisfies

  • After rematerializing a tensor, one additional eviction can at most double the gap to 2L, and that bound is maintained across the remaining gradients.The argument applies to every gradient in the sequence.Source: §2 · p. 19

B PROOF OF THEOREM 3.2

The proof constructs an adversarial graph one node at a time in response to deterministic DTR’s choices. On that graph, DTR can require at least an Ω(N/B) runtime factor over static checkpointing.

B PROOF OF THEOREM 3.2

  • The adversary reveals each next node based on DTR’s prior choices, while DTR can use only the graph revealed so far.This construction targets deterministic heuristics.Source: B PROOF OF THEOREM 3.2 · p. 20
  • With B paths descending from t0 but only B−1 descendant memory units, some path has no resident tensors; extending it forces DTR to rematerialize that path.The adversary repeats this step as it reveals the graph.Source: B PROOF OF THEOREM 3.2 · p. 20
  • The static algorithm can compute one path and its children at a time using two memory units, reusing them across paths without rematerialization.Its total cost is Θ(N).Source: B PROOF OF THEOREM 3.2 · p. 21
  • DTR requires Ω(N²/B) tensor computations on the adversarial graph, versus Θ(N) for static checkpointing, a runtime factor of at least Ω(N/B).The static algorithm computes the graph one path at a time.Source: B PROOF OF THEOREM 3.2 · p. 21

C SIMULATOR SPECIFICATION

The simulator models DTR over tensor storages, tracks metadata and dependencies to guide eviction and rematerialization, and replays logged operations. Its specification also defines heuristic variants and runtime policies.

C.1 FUNDAMENTAL ABSTRACTIONS

  • DTR models memory through storages and tensors as views of buffers, tracking aliases, references, residency, locks, and whether storages are evictable.Operators are modeled as pure functions with compute costs.Source: §C.1 · pp. 21–22

C.2 FORMAL METADATA DEFINITIONS

  • The evicted neighborhood estimates which evicted storages are needed to recompute a storage or depend on it; an undirected Union-Find approximation tracks component costs.This avoids the complexity of exact dynamic graph tracking, but rematerialization can leave phantom connections in the approximation.Source: §C.2 · p. 23

C.3 FORMAL HEURISTIC DEFINITIONS

  • hDTR scores candidate storages using compute cost, size, and staleness, evicting the lowest-scored storage.The heuristics include exact, approximate, and local cost variants.Source: §C.3 · pp. 23–24

C.4 IMPLEMENTATION DETAILS

  • The simulator tracks a memory budget, current usage, time, heuristic scores, and evictable storages; rematerialization locks parents and recursively restores missing dependencies.Evictions reduce memory use, while rematerialization accounts for operator costs and output storage sizes.Source: §C.4 · p. 24

C.5 ADDITIONAL RUNTIME OPTIMIZATIONS

  • When the source program deallocates a storage, DTR can banish it or eagerly evict it; eager eviction avoids over-pinning memory but cannot evict constants.In practice, eager eviction supported lower budgets by pinning fewer values.Source: §C.5 · pp. 24–25

C.6 LOG-REPLAYING MECHANISM

  • The simulator replays logged tensor operations and reference-count events, and models in-place mutation by cloning changed inputs and updating their mappings.Live tensors at log end are rematerialized and locked as necessary outputs.Source: §C.6 · pp. 25–26

D ABLATION STUDY

  • The ablation study compares metadata choices and deallocation policies using the same model logs as the main evaluation.It measures both rematerialization work and heuristic or metadata-maintenance storage accesses.Source: D ABLATION STUDY · p. 26

D.1 DATA SOURCES

The ablation varies staleness and size metadata while fixing compute-cost measures, then compares the resulting h′DTR variants. More precise cost information generally improves savings, but the useful metadata depends on model architecture.

D.1 DATA SOURCES

  • The ablation compares h′DTR scores c(t)/[m(t) · s(t)] while varying whether staleness and size are included and fixing c to e∗, EqClass, local, or no cost.The resulting non-ablated variants use e∗, approximate EqClass, and local cost, respectively.Source: §D.1 · pp. 26–27
  • More precise evicted-neighborhood costs generally enable greater savings, while staleness and size are needed for acceptable computational overhead.The approximate EqClass cost performs comparably to exact e∗ while requiring less information.Source: §D.1 · p. 26, Figure 7 · p. 26, Figure 8 · p. 27
  • For static DenseNet, ResNet, and UNet, cost or size alone outperforms cost plus staleness; for dynamic models, the reverse pattern holds.The paper notes that architecture-related explanations remain uncertain and merit further research.Source: §D.1 · p. 26

D.2 BANISHING AND DEALLOCATIONS

Eager eviction supports lower memory budgets than banishing across most tested models, while both deallocation-aware policies reduce rematerialization overhead relative to ignoring deallocations. The appendix describes runtime overhead measurements, the PyTorch integration, and optimization tradeoffs.

D.2 BANISHING AND DEALLOCATIONS

  • Eager eviction supports lower budgets than banishing across most tested models; on UNet, it supports 50% of baseline budget, while banishing supports 90% and OOMs at 0.8 ratio.Banishing can nevertheless achieve lower computational overhead under the same budget and savings for ResNet.Source: §D.2 · p. 28, Figure 11 · p. 28
  • Compared with ignoring deallocations, both banishing and eager eviction yield noticeably lower rematerialization overhead.The runtime-overhead experiment counts storage accesses for heuristic evaluations and metadata maintenance rather than relying on simulator wall-clock time.Source: §D.2 · p. 28, §D.3 · p. 28

D.3 RUNTIME OVERHEAD

  • Heuristic storage accesses are generally separated by at least an order of magnitude, and DTR is presented as a middle ground in runtime and computational overhead.The implementation recomputes heuristics often, and optimized score maintenance could reduce overhead at the cost of greater implementation complexity.Source: §D.3 · pp. 28–29, Figure 12 · p. 28

E.1 INTEGRATION INTO PYTORCH

  • The PyTorch prototype wraps tensors as CheckpointTensors, dispatches operations through overloads, and tracks references, aliases, mutations, and multi-output results.Its mutation layer uses copy-on-write; alias pools are evicted together, while multi-output tensors can be evicted separately and rematerialized together.Source: §E.1 · pp. 29–30

E.2 RUNTIME OPTIMIZATIONS

  • Runtime search overhead comes from recomputing staleness and equivalence-class costs; approximate filters skip small tensors and can sample √n of n evictable tensors.Random sampling sometimes failed at low budgets or very large inputs by excluding good eviction candidates; deeper PyTorch changes could reduce other overhead sources.Source: §E.2 · p. 30

E.3 HANDLING ERRORS IN TRIALS

The DTR prototype had out-of-memory errors and occasional hangs on some low-budget or large-input trials; disabling random sampling resolved most OOMs but not all UNet failures. The authors describe a retry procedure for large-input UNet measurements and discuss timing variability as a possible source of failures.

E.3 HANDLING ERRORS IN TRIALS

  • Most CUDA out-of-memory errors disappeared when random sampling was disabled, but the causes of hangs remained uncertain and they were harder to reproduce.The authors could not determine whether hangs came from deep recursive rematerialization or a PyTorch loop or deadlock.Source: §E.3 · pp. 30–31
  • UNet’s largest batch sizes still had failures after sampling was disabled; the reported runs had 1 failure at batch size 9 and 10 at batch size 10.The authors retried failed runs, defining a hang as a trial taking twice as long as the preceding trial; most failures at batch size 10 occurred during warm-up.Source: §E.3 · p. 31
  • Variance in operator timings may affect rematerialization decisions and contribute to occasional UNet failures.The authors suggest fixed operator cost models or proxy measures, such as operation-count staleness, to reduce timing variation.Source: §E.3 · p. 31
Loading 2006.09616v4…