Libra: A Python-Level Tensor Re-Materialization Strategy for Reducing Deep Learning GPU Memory Usage

To tackle more complex and challenging problems, Deep Neural Networks (DNNs) are becoming larger and deeper, leading to exponential growth in model sizes. In contrast, on-device memory grows linearly, making limited memory resources a bottleneck that restricts the development of advancing DNN architectures. In this paper, we propose Libra, a deep learning framework that uses a greedy refinement approach to generate a memory optimization re-materialization plan through swapping and recomputation to reduce peak memory usage of DNN training. To address the limitations of swapping or recomputation, Libra employs a hybrid policy, setting recomputation as the initial policy due to its low re-generation time and using swapping as a trade-off to increase training time overhead but significantly reduce the memory overhead associated with the recomputation chain that will create a new peak. Evaluations on various DNN models shows that Libra achieves exceptional memory reduction performance, increasing the maximum batch size by 3.12x compared to PyTorch, and by 1.42x and 2.57x in average compared to other state-of-the-art baselines.