Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory
This work addresses memory limitations in training deep neural networks, which is a critical bottleneck for researchers and practitioners, though it is incremental as it builds on existing checkpointing techniques.
The paper tackles the problem of high memory usage in training deep neural networks by introducing a new activation checkpointing method that dynamically selects which forward activations to save, enabling significant memory reduction. The result is a PyTorch implementation that consistently outperforms existing approaches across various networks, image sizes, and batch sizes.
This paper introduces a new activation checkpointing method which allows to significantly decrease memory usage when training Deep Neural Networks with the back-propagation algorithm. Similarly to checkpoint-ing techniques coming from the literature on Automatic Differentiation, it consists in dynamically selecting the forward activations that are saved during the training phase, and then automatically recomputing missing activations from those previously recorded. We propose an original computation model that combines two types of activation savings: either only storing the layer inputs, or recording the complete history of operations that produced the outputs (this uses more memory, but requires fewer recomputations in the backward phase), and we provide an algorithm to compute the optimal computation sequence for this model. This paper also describes a PyTorch implementation that processes the entire chain, dealing with any sequential DNN whose internal layers may be arbitrarily complex and automatically executing it according to the optimal checkpointing strategy computed given a memory limit. Through extensive experiments, we show that our implementation consistently outperforms existing checkpoint-ing approaches for a large class of networks, image sizes and batch sizes.