Overview of DeepSpeed-TED: Hybrid Parallelism in MoE Models
The paper explores the challenges and solutions associated with training Mixture-of-Experts (MoE) models, which have gained traction due to their ability to increase parameter count without proportionate computational costs. While traditional frameworks fall short when accommodating large base models in MoE structures, the authors present a novel approach—DeepSpeed-TED. This work integrates three dimensions of parallelism: tensor, expert, and data, thereby facilitating training on significantly larger base models compared to prior solutions.
Core Contributions
- Three-Dimensional Hybrid Parallelism: The paper introduces DeepSpeed-TED, which combines tensor parallelism from Megatron-LM, expert parallelism from DeepSpeed-MoE, and data parallelism from ZeRO. This integration permits the training of MoE models with base models that are 4-8 times larger than those supported by current state-of-the-art frameworks.
- Memory and Communication Optimizations: The framework addresses two critical bottlenecks:
- Memory Usage: By implementing a tiled optimizer, the solution alleviates memory spikes during the optimizer step. This allows models to train successfully even with a substantial increase in the number of experts and base model size.
- Communication Overhead: To mitigate increased communication costs, the authors propose Duplicate Token Dropping (DTD) and Communication-aware Activation Checkpointing (CAC). These optimizations reduce unnecessary data movement and lower collective communication times, respectively.
- Empirical Evaluation: The framework's efficacy is underscored through extensive experiments on Summit and ThetaGPU supercomputers. Results indicate significant improvements in hardware efficiency and training speed, particularly when applying DTD and CAC.
Technical Approach
The paper elaborates on distinct dimensions of parallelism being used for different components of MoE models. Notably:
- Tensor Parallelism: Partitioning computation at a layer level across GPUs to allow large base models to be manageable.
- Expert Parallelism: Allocating experts across GPUs in an embarrassingly parallel manner.
- Data Parallelism: Utilizing ZeRO to efficiently shard optimizer states, further decreasing memory consumption.
The new partitioning approach resolves two observed bottlenecks—inefficient memory usage during optimizer operations and excessive communication across GPUs—both exacerbated by increasing expert counts and base model sizes.
Results and Implications
The paper demonstrates that DeepSpeed-TED supports MoE models with several billion parameters, achieving significant efficiency improvements:
- Speedups of up to 29% in communication-intensive operations.
- Supporting MoE models substantially larger than those previously feasible.
Such advancements imply that larger, highly performant MoE models can now be trained without prohibitive scaling of computational resources. This has practical implications for the scaling laws of neural networks, enabling further exploration into model architecture design and efficiency.
Future Directions
In light of these enhancements, future research could explore integrating pipeline parallelism to extend the framework's scalability. Additionally, there lies potential in further optimizing communication routines for diverse network architectures and decreasing resting state memory requirements using novel data redundancy techniques.
DeepSpeed-TED represents a significant step in distributed deep learning, offering a robust solution to the complexity of training massive MoE models while maintaining efficiency in computational and memory resources necessary for operations on supercomputing infrastructures.