TransAware: An Automatic Parallel Method for Deep Learning Model Training with Global Model Structure Awareness
摘要
The distributed training strategy for deep learning models is key to the development and industrialization of large models, such as Bert, GPT-2, and Llama2. However, it is challenging to find an optimal distributed training strategy because of the complexity of the model structure and the large scale of parameters. The automatic parallel method based on reinforcement learning is a mainstream method to automatically find distributed training strategies for models. The existing methods mainly focus on the feature extraction of the local structure of models, and lack the perception of the global structure of models, resulting in poor performance of the distributed training strategy generated by reinforcement learning. To address the problem, this paper proposes an automatic parallel method for deep learning model training with global model structure awareness, TransAware. It uses a globally-aware graph embedding method to exact global model structure, then it makes full use of the self-learning ability of the reinforcement learning method to iteratively search for the optimal distributed training strategy for models. In particular, a dual-layer node fusion method to reduce the search space of distributed training strategy. The experiments show that the model training performance with the distributed training strategy searched by TransAware can be improved by up to 12.82%, 15%, and 10.53%, compared with Hierarchical, Placeto, and Aware, respectively.