MoNTA: Accelerating Mixture-of-Experts Training with Network-Traffc-Aware Parallel Optimization
Year: 2024
Type: preprint
Abstract: The Mixture of Experts (MoE) is an advanced model architecture in the industry that combines multiple specialized expert models from various domains into a single supermodel. This approach enables the model to scale without significantly increasing the computational costs of training and inference, while maximizing model performance. However, current distributed training frameworks do not consider the ultimate optimization of communication, especially for large base models. This paper proposes a network-traffic-aware parallel optimization method that selects the optimal parallel strategy based on the communication volume, and the training cluster's inter-node and intra-node network topologies. Compared to the DeepSpeed, MoNTA achieves an 8x increase in AllToAll communication performance under 8-card tensor parallelism. Compared to the baseline, training a 2x70B model using 16 A800 cards, with an 8K sequence, results in a 13% overall latency performance improvement. Project Page: https://github.com/EnflameTechnology/DeepSpeed.
Source: arXiv (Cornell University)
Authors: Guo, Jingming, Liu, Yan, Meng, Yu, Tao, Zhiwei, Liu, Banglan, Chen, Gang, Li, Xiang
Institution: Enflame (China)
Language: English
Subfield: Hardware and Architecture
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
OpenAlex ID: W4404345293