Dynamic Tensor RematerializationDownload PDF

28 Sep 2020 (modified: 14 Jan 2021)ICLR 2021 SpotlightReaders: Everyone
  • Keywords: Rematerialization, Memory-saving, Runtime Systems, Checkpointing
  • Abstract: Checkpointing enables training deep learning models under restricted memory budgets by freeing intermediate activations from memory and recomputing them on demand. Previous 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(\smash{\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 just by interposing on tensor allocations and operator calls and collecting lightweight metadata on tensors.
  • One-sentence Summary: We present an online algorithm for rematerialization (recomputing intermediate activations during backpropagation instead of storing them), which enables training under low memory, finding that it is competitive with offline techniques.
  • Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics
  • Supplementary Material: zip
11 Replies