Dynamic Tensor Rematerialization

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.

arXiv comment31 pages, 12 figures, implementation available here: https://github.com/uwsampl/dtr-prototype, OpenReview: https://openreview.net/forum?id=Vfs_2RnOD0H

References in corpus (8)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Large Scale GAN Training for High Fidelity Natural Image Synthesis
• Training Deep Nets with Sublinear Memory Cost
• Reformer: The Efficient Transformer
• The Reversible Residual Network: Backpropagation Without Storing Activations
• Glow: Graph Lowering Compiler Techniques for Neural Networks
• Low-Memory Neural Network Training: A Technical Report
• Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory

Cited by in corpus (1)

• ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training