MemOpt: a memory optimization method for deep learning model training based on dual intelligent reinforcement learning

Abstract In recent years, with the rapid growth in the scale of datasets and neural network models, there has been a significant imbalance between the memory requirements during model training and the memory resources available on training devices. Existing memory optimization techniques like recomputation, memory swapping, and their adaptive combinations do not fully consider the structural information of the model and overlook the impact of application costs and timing on training efficiency. Addressing this issue, this paper proposes a memory optimization method called MemOpt, which uses dual-agent reinforcement learning to dynamically search for appropriate memory optimization strategies and execution timing based on model structure and device information. It can optimize memory without compromising accuracy while minimizing additional overhead. Experimental results show that the MemOpt method significantly increases the maximum batch size for model training by up to 8.7% and training throughput by up to 42.3% compared with baseline methods. In the future, this method may find better applications in large-scale neural networks.