Training modern AI models requires managing billions of parameters with vast amounts of data, typically on GPU devices that face limitations due to limited internal memory. This can complicate or even hinder the training of large foundation models. To address this, we utilize a Directed Acyclic Graph (DAG) model, which effectively represents both traditional memory management tasks and novel neural networks. DAGs enable clear visualization and optimization of data and computation flow across various layers of neural networks. Our objective is to determine an optimal topological order that minimizes memory peak usage. Traditional approaches, such as Depth First Search(DFS) and greedy algorithms, yield suboptimal results, while conventional dynamic programming (DP) methods have high time complexity. In this paper, we enhance the traditional DP algorithm and apply our new method to three randomly wired neural networks. Our empirical results show that our method outperforms existing algorithms in most cases and significantly reduces the time cost of the original DP method, with minimal impact on solution quality. Consequently, our solution shows great promise for large model training applications.