Dynamic Tensor Rematerialization

As state-of-the-art deep learning (DL) models continue to grow, training them within the constraints of on-device memory becomes increasingly challenging. The memory demands of recent models prevent their training on memorylimited devices (such as specialized accelerators, low-powered embedded devices, or older GPUs) and limit researchers' ability to explore memory-intensive architectures or training techniques. Checkpointing enables training with models and batches that are larger than would otherwise fit in on-device memory without modifying the model's design. This is achieved by freeing some activations from memory and recomputing them on demand. Adapted from techniques in automatic differentiation [Baydin et al., 2015, Griewank and Walther, 2000, Siskind and Pearlmutter, 2018], 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.