DOME: Dynamic Optimization of GPU Memory for Training DNN Models with Limited Resources

To improve data processing and learning capabilities, DNN models are developing towards increasingly complex structures and a huge number of parameters. However, the limited memory capacity of specialized acceleration chips such as GPUs has become a critical constraint in exploring advanced DNN models. Memory swapping and recomputation are two effective GPU memory optimization techniques that enable the training of larger models in memory-constrained hardware environments. Existing methods mostly require prior knowledge of the complete computational graph, which is not suitable for the dynamic computational graph mode adopted by mainstream deep learning frameworks. This paper proposes DOME, a tensor-level GPU memory optimization method for dynamic computational graphs. DOME combines tensor swapping and recomputation, inspired by cache replacement algorithms, to release the least costly tensors when GPU memory is insufficient, thereby supporting larger batch sizes during model training. DOME employs a heuristic evaluation function with higher training efficiency, estimating the cost of tensor swapping or recomputation using information obtained online at runtime. Experiments demonstrate that, for different DNN models, DOME can increase the maximum training batch size by 1.56 to 2.33 times. Additionally, DOME can improve training throughput by 9.8% to 35.0% compared to SOTA methods.