What is EPLB
EPLB (Expert Parallelism Load Balancer) is DeepSeek 's expert parallelism load balancer, designed to solve the problem of uneven load across different expert models (Experts) in large-scale model training. EPLB is based on a redundant expert strategy, replicating high-load experts and reasonably distributing them across different GPUs to achieve load balancing. EPLB combines group-limited expert routing technology to place experts in the same group within the same node, reducing cross-node communication overhead. EPLB offers two load balancing strategies: Hierarchical Load Balancing and Global Load Balancing, suitable for different scenarios. By optimizing the replication and placement of expert models, EPLB can significantly improve GPU resource utilization and training efficiency.
EPLB's Main Features
- Load Balancing: Based on the estimated load of experts (Experts), dynamically adjusts the expert replication and allocation strategy to ensure that the load difference between different GPUs is minimized.
- Expert Replication: Based on a redundant expert strategy, replicates high-load experts to alleviate the problem of uneven load.
- Resource Optimization: Maximizes the utilization of GPU resources, reduces performance bottlenecks caused by uneven load, and improves model training efficiency.
- Communication Optimization: A reasonable expert placement strategy reduces inter-node communication overhead and lowers communication latency.
- Flexible Strategy Support: Provides two strategies, Hierarchical Load Balancing and Global Load Balancing, suitable for different scenarios and stages.
- Multi-Layer MoE Model Support: Suitable for multi-layer Mixture of Experts (MoE) models, handles complex model structures, and supports flexible expert allocation and mapping.
EPLB's Technical Principles
- Redundant Expert Strategy: In expert parallelism, the load of different experts may vary due to input data and model structure. Introduce redundant experts (replicating high-load experts) to balance the load. Support high-load experts being replicated multiple times and distributed across multiple GPUs to avoid overloading a single GPU.
- Hierarchical Load Balancing: Evenly distribute expert groups across different nodes to ensure that the load of each node is roughly equal. Within each node, further replicate experts and distribute the replicated experts to GPUs within the node to ensure intra-node load balancing. Try to place experts in the same group within the same node to reduce cross-node communication overhead.
- Global Load Balancing: In other cases (such as when the number of nodes cannot evenly divide the number of expert groups or when larger-scale parallelism is required), adopt a global strategy: ignore the restriction of expert groups and globally replicate and distribute experts across all available GPUs. Based on dynamically adjusting the number of expert replicas and placement positions, ensure global load balancing.
- Load Estimation and Dynamic Adjustment: EPLB relies on estimated values of expert load to guide the load balancing strategy. Load estimation is based on historical statistical data (such as moving averages). Based on the load estimates, dynamically adjust the expert replication and allocation strategy to adapt to different training stages and data distributions.
- Expert Mapping and Resource Allocation: Based on the rebalance_experts function, output the expert replication and placement plan, and map experts to specific GPUs. The output mapping relationships include physical-to-logical (phy2log) and logical-to-physical (log2phy) mappings, as well as the number of replicas for each expert (logcnt).
EPLB project address
- GitHub repository:https://github.com/deepseek-ai/eplb
Core modes of EPLB
- Hierarchical Load Balancing mode (Hierarchical Load Balancing): When the number of nodes can evenly divide the number of expert groups, based on a hierarchical load balancing strategy, optimize the load distribution within and across nodes.
- Global Load Balancing mode (Global Load Balancing): When the number of nodes cannot evenly divide the number of expert groups or larger-scale parallelism is required, based on global replication and expert distribution, achieve overall load balancing.
EPLB code demonstration example
- The code demonstrates an example of a two-layer MoE model, with each layer containing 12 experts. Each layer introduces 4 redundant experts, with a total of 16 replicas placed on 2 nodes, each containing 4 GPUs.
Application scenarios of EPLB
- Large-scale distributed training: Suitable for multi-node, multi-GPU environments, based on flexibly switching between hierarchical or global load balancing modes, optimize resource utilization and reduce communication overhead.
- Prefill stage: In the early stage of model training, based on hierarchical load balancing, reduce cross-node communication and improve small-scale parallelism efficiency.
- Decode stage: In the later stage of training when large-scale parallelism is required, use global load balancing to dynamically adjust the load and handle complex tasks.
- Heterogeneous hardware environment: When the number of nodes does not match the number of expert groups, the global load balancing mode can flexibly adapt to heterogeneous configurations and achieve efficient load balancing.
- Dynamic load changes: For dynamic changes in load during the training process, combine hierarchical or global load balancing strategies to adjust in real time, ensuring the efficiency and stability of the training process.