Expert Parallelism Load Balancer (EPLB)

In the world of AI and deep learning, managing the computational load efficiently is a critical task. Models, especially large-scale ones like those used in Natural Language Processing (NLP) or computer vision, require handling massive amounts of data and computation across multiple computing resources such as GPUs or server nodes. This task often involves parallelism, which is the strategy of dividing the workload into smaller chunks and processing them simultaneously. Recently DeepSeek open sourced their EPLB algorithm : https://github.com/deepseek-ai/EPLB

A central concept in parallelism, especially for models with large parameters (such as mixture-of-experts or MoE models), is expert parallelism, where we distribute expert models (or parts of a model) across different resources. This helps to balance the computational load across multiple GPUs, nodes, or physical experts while optimizing performance and minimizing bottlenecks.

Let's explore the code snippet that handles expert parallelism load balancing through replication and rebalancing in a hierarchical structure.

1. Packing Objects with Balanced Weights

The first step in this parallelism process involves packing objects (data or model components) into groups (or "packs"). This ensures that each group has an approximately equal weight distribution, minimizing the chance of overloading any one pack.

Function: balanced_packing

This function's goal is to divide n weighted objects into m groups or "packs" such that the distribution of weight in each pack is as balanced as possible. Here's a breakdown of its core logic:

• Input: A tensor weight of size [X, n], where X is the number of layers, and n is the number of objects to pack. Additionally, num_packs specifies how many packs or groups to create.

• Output:

  • pack_index: Tensor showing which pack each object belongs to.

  • rank_in_pack: Tensor showing the rank (or position) of the item within its respective pack.

The function first checks if the number of groups divides evenly into the number of packs. If this is true, it proceeds to sort the objects by weight and assigns them to packs, ensuring that each pack's weight is as balanced as possible.

2. Replicating Experts to Minimize Load

Once the objects are packed, the next step is to replicate these "logical experts" (think of these as model components or neurons) into physical experts. The goal here is to distribute the computational load evenly across physical devices like GPUs.

Function: replicate_experts

This function replicates logical experts (model components) into physical experts, minimizing the load imbalance across the physical replicas:

• Input:

  • weight: A tensor representing the weight of each logical expert.

  • num_phy: The total number of physical experts (often corresponds to available GPUs).

• Output:

  • phy2log: Mapping from physical expert IDs to logical expert IDs.

  • rank: The rank (or order) of each replica within its physical expert.

  • logcnt: Number of replicas each logical expert has across physical experts.

The idea here is that when we replicate a logical expert across multiple physical devices, we aim to minimize the maximum load (or number of tasks) that each physical expert carries. By distributing the work evenly, we ensure better resource utilization and prevent some devices from being overloaded while others are underutilized.

3. Hierarchical Rebalancing for Optimal Distribution

In more complex scenarios where multiple nodes or GPUs are involved, we need to rethink the way experts are distributed across devices. A hierarchical structure allows for the efficient distribution of logical experts across various levels, starting from nodes down to GPUs, ensuring optimal use of available resources.

Function: rebalance_experts_hierarchical

This function deals with rebalancing experts in a hierarchical fashion. It distributes logical experts across multiple server nodes, and within each node, it ensures that they are evenly distributed across GPUs:

• Input:

  • weight: The load statistics for all logical experts.

  • Other parameters define the hardware layout: num_physical_experts, num_groups, num_nodes, num_gpus.

• Output:

  • physical_to_logical_map: A mapping from physical experts to logical experts.

  • logical_to_physical_map: A mapping from logical experts to physical experts.

  • logical_count: The count of replicas per logical expert.

4. Main Function: Rebalance Experts

The main function, rebalance_experts, serves as the entry point for balancing the load across multiple replicas, groups, nodes, and GPUs:

• Input: Similar to the hierarchical function, this function takes a tensor of weights for logical experts and distributes them across multiple replicas, nodes, and GPUs.

• Output:

  • physical_to_logical_map: Mapping physical experts to logical ones.

  • logical_to_physical_map: Detailed mapping for each logical expert to all its corresponding physical replicas.

  • logcnt: Tracks the number of replicas for each logical expert.

Why Is This Approach Useful?

This hierarchical load-balancing strategy is essential when dealing with large-scale models, particularly in mixture-of-experts (MoE) models, where only a subset of experts are active at any given time.

1. Efficient Resource Usage: By evenly distributing the computational load, this method ensures no single GPU or node becomes a bottleneck.

2. Scalability: As the number of nodes, GPUs, or experts grows, the system can scale without introducing inefficiencies or resource underutilization.

3. Minimal Load Imbalance: Using AI-based strategies like packing, ranking, and replicating ensures that the maximum load per physical expert is minimized, improving overall training and inference performance.

Conclusion

Efficient expert parallelism is a cornerstone of scaling AI models across large clusters of GPUs and nodes. This load-balancing strategy, incorporating techniques like balanced packing and hierarchical expert distribution, ensures that computational resources are used optimally, preventing resource contention and ensuring fast and efficient model training. The implementation described is just one approach to managing expert parallelism, but it highlights the importance of balancing the workload in distributed AI systems, especially as model sizes continue to grow.

By applying these methods, AI engineers can take full advantage of multi-GPU and multi-node environments, ensuring that large models can be trained efficiently and at scale.