EPLB - DeepSeek's open-source expert parallel load balancer

What is EPLB?

EPLB (Expert Parallelism Load Balancer) is an expert parallel load balancer introduced by DeepSeek, which solves the problem of uneven load among different expert models during large-scale model training. EPLB uses a redundant expert strategy, replicating high-load experts and distributing them rationally across different GPUs to achieve load balancing. EPLB incorporates group-limited expert routing technology, placing 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.

Main functions of EPLB

• Load balancingBased on the load estimates from experts, dynamically adjust the experts' replication and allocation strategies to ensure that load differences between different GPUs are minimized.
• Expert CopyBased on the redundancy expert strategy, high-load experts are replicated to alleviate the problem of load imbalance.
• Resource optimizationMaximize the use of GPU resources, reduce performance bottlenecks caused by uneven load, and improve model training efficiency.
• Communication optimizationA reasonable expert placement strategy reduces communication overhead between nodes and lowers communication latency.
• Flexible strategy supportIt provides two strategies: Hierarchical Load Balancing and Global Load Balancing, suitable for different scenarios and stages.
• Multilayer MoE model supportSuitable for multi-layered hybrid expert models (MoE), handling complex model structures, and supporting flexible expert assignment and mapping.

EPLB Technical Principles

• Redundant expert strategyIn expert parallelism, the workload of different experts may vary depending on the input data and model structure. Introducing redundant experts (replicating high-load experts) balances the load. It supports replicating high-load experts multiple times and distributing them across multiple GPUs, avoiding overloading a single GPU.
• Hierarchical load balancingDistribute expert groups evenly across different nodes to ensure that the load on each node is roughly equal. Within each node, further replicate the experts and distribute the replicated experts across the node's GPUs to ensure load balancing within the node. Try to place experts from the same group on the same node to reduce cross-node communication overhead.
• Global load balancingIn other cases (such as when the number of nodes is not divisible by the number of expert groups or when larger-scale parallelism is required), a global strategy is adopted: ignoring the limitations of expert groups, experts are globally replicated and distributed across all available GPUs. Global load balancing is ensured by dynamically adjusting the number of expert replicas and their placement.
• Load estimation and dynamic adjustmentEPLB relies on estimates of expert load to guide its load balancing strategy. Load estimates are based on historical statistics (such as moving averages). Based on these load estimates, the replication and allocation strategies for experts are dynamically adjusted to adapt to different training phases and data distributions.
• Expert Mapping and Resource AllocationBased on the `rebalance_experts` function, the replication and placement plan of experts is output, mapping experts to specific GPUs. The output mapping relationships include physical to logical (phy2log) and logical to physical (log2phy) mappings, as well as the replication count (logcnt) for each expert.

EPLB project address

• GitHub repository:https://github.com/deepseek-ai/eplb

EPLB's core mode

• Hierarchical Load BalancingWhen the number of nodes is divisible by the number of expert groups, a hierarchical load balancing strategy is used to optimize load distribution within and between nodes.
• Global Load BalancingWhen the number of nodes cannot be divided by the number of expert groups or when larger-scale parallelism is required, overall load balancing is achieved based on global replication and expert allocation.

EPLB code demonstration example

• The code demonstrates an example of a two-layer MoE model, with 12 experts per layer. Four redundant experts are introduced per layer, for a total of 16 replicas placed on two nodes, each node containing four GPUs.

Application scenarios of EPLB

• Large-scale distributed trainingSuitable for multi-node, multi-GPU environments, it optimizes resource utilization and reduces communication overhead by flexibly switching between hierarchical or global load balancing modes.
• Prefilling stageIn the early stages of model training, hierarchical load balancing reduces cross-node communication and improves small-scale parallel efficiency.
• Decoding stageWhen large-scale parallelism is required in the later stages of training, global load balancing is used to dynamically adjust the load to cope with complex tasks.
• Heterogeneous hardware environmentWhen 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 changesTo address the dynamic changes in load during training, a hierarchical or global load balancing strategy is used to adjust in real time, ensuring the efficiency and stability of the training process.