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.

Mixture-of-Experts (MoE) efficiently trains large models by using sparse activation to lower costs, selecting a few experts based on data characteristics. However, it faces challenges such as All-to-All communication overhead and load imbalance, with most optimizations targeting dynamic graphs rather than the more efficient static graphs. This study identifies two key challenges in training MoE on static graphs: 1) excessive All-to-All communication (up to 75% of iteration time) and load imbalance (70% of tokens handled by two experts) between experts due to the sparse structure of the MoE model and the token distribution; and 2) inefficient zero-padding for static shapes, leading to unnecessary computational overhead(wasting approximately 50% of resources). Thus, EfficientMoE, a scheduling method based on expert load and data characteristics, is introduced. EfficientMoE first designs a sampler to collect real-time information about token distribution, expert load, etc. It constructs a load prediction model to evaluate expert load. Subsequently, EfficientMoE proposes a dynamic schedule strategy for experts with evaluated expert load, reducing All-to-All communication and addressing load-balancing issues. Additionally, an expert capacity model is proposed to set different capacities for replicas of hot experts before static graph compilation, minimizing computation and storage overhead caused by significant padding. This study implements EfficientMoE in MindSpore and uses 32 Ascend AI accelerators to train an MoE model with 21 billion parameters and evaluate its validity. EfficientMoE demonstrated an improvement of 30% in model training time, approximately 12% reduction in communication time, and saved 35% computational resources across different clusters, compared with Switch transformers, and the Fastermoe method for static graphs.

As the size of datasets and neural network models increases, automatic parallelization methods for models have become a research hotspot in recent years. The existing auto-parallel methods based on machine learning or graph algorithms still have issues with search efficiency and applicability. This paper proposes an automatic parallel method based on a dual-population genetic algorithm, TGA, which transforms model partitioning and placement into an integer linear programming problem and constructs a cost model to evaluate the solution. The solution space is built using the neural network’s dataflow graph and device cluster’s topology, and the dual-population genetic algorithm is used to search for the optimal model parallel strategy. Experiments with various models show that the proposed method can improve single-step execution time by up to 42% compared to the Baechi method and up to 37.7% compared to the Hierarchical method.