Dynamic Tensor Rematerialization

By Marisa Kirisame, Steven Lyubomirsky, A. Haan, Jennifer Brennan, Mike He, Jared Roesch, Tianqi Chen, Zachary Tatlock. 2020

119 citations recorded by the source corpus.

Checkpointing enables training larger models by freeing intermediate activations and recomputing them on demand. Previous checkpointing techniques are difficult to generalize to dynamic models because they statically plan recomputations offline. We present Dynamic Tensor Rematerialization (DTR), a greedy online algorithm for heuristically checkpointing arbitrary models. DTR is extensible and general: it is parameterized by an eviction policy and only collects lightweight metadata on tensors and operators. Though DTR has no advance knowledge of the model or training task, we prove it can train an $N$-layer feedforward network on an $\Omega(\sqrt{N})$ memory budget with only $\mathcal{O}(N)$ tensor operations. Moreover, we identify a general eviction heuristic and show how it allows DTR to automatically provide favorable checkpointing performance across a variety of models and memory budgets.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Hardware and Architecture

In Efficiency & Compression

Citations in the corpus

PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, background): “Orthogonal memory-saving techniques include gradient compression [1], mixed-precision training [6], tensor rematerialization [11] and CPU-offloading [3], but they could have implications on model accuracy and incur overhead in (un)compression, quantization, recomputation, and host-to-device copies, ”
Software-hardware co-design for fast and scalable training of deep learning recommendation models (2021, background): “DeepSpeed [50] fully shardsmodel parameters, gradients and optimizer states across all nodes, and reconstructs necessary states on the fly using checkpoint partitioning and rematerialization [21, 28] to drastically reduce memory usage.”
Wavelet: Efficient DNN Training with Tick-Tock Scheduling (2021, methodology): “To further reduce the memory overheads for holding intermediate results, we can also leverage techniques like tensor re-materialization (Kirisame et al., 2020; Jain et al., 2020) to recompute intermediate results on demand (Huang et al., 2019).”
Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging (2022, methodology): “Prior work has proposed strategies including paging to auxiliary memory (Peng et al., 2020) and rematerialization (Chen et al., 2016; Jain et al., 2020; Kirisame et al., 2021) to reduce the memory footprint of training in the cloud.”
POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging (2022, methodology): “Prior work has proposed strategies including paging to auxiliary memory (Peng et al., 2020) and rematerialization (Chen et al., 2016; Jain et al., 2020; Kirisame et al., 2021) to reduce the memory footprint of training in the cloud.”
High-performance, Distributed Training of Large-scale Deep Learning Recommendation Models (2021, background): “DeepSpeed [44] fully shards model parameters, gradients and optimizer states across all nodes, and reconstructs necessary states on the fly using checkpoint partitioning and rematerialization [17, 24] to drastically reduce memory usage.”
This paper is included in the Proceedings of the 16th USENIX Symposium on Operating Systems Design and Implementation. (unknown): “In addition to parallelization, there are other complementary techniques for training large-scale models, such as memory optimization [12, 14,21,23,28,46], communication compression [6,53], and low-precision training [35].”