MoNTA: Accelerating Mixture-of-Experts Training with Network-Traffc-Aware Parallel Optimization
This addresses a specific bottleneck for practitioners training large MoE models, representing an incremental optimization.
The paper tackles the communication bottleneck in distributed training of Mixture-of-Experts models by proposing a network-traffic-aware parallel optimization method, achieving an 8x increase in AllToAll communication performance and a 13% overall latency improvement for a 2x70B model.
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.