专家并行负载均衡器 (EPLB)

在使用专家并行（EP）时，不同的专家会被分配到不同的 GPU 上。由于不同专家的负载可能因当前工作负载而异，保持各 GPU 的负载均衡至关重要。正如 DeepSeek-V3 论文中所述，我们采用了一种冗余专家策略，即对负载较高的专家进行复制。随后，我们通过启发式方法将这些复制的专家合理地分配到各个 GPU 上，以确保不同 GPU 之间的负载均衡。此外，得益于 DeepSeek-V3 中使用的分组受限专家路由机制，我们还会尽可能地将同一分组内的专家放置在同一节点上，从而减少节点间的数据传输开销。

为了便于复现和部署，我们将已部署的 EP 负载均衡算法开源至 eplb.py 文件中。该算法会根据预估的专家负载计算出一个均衡的专家复制与放置方案。需要注意的是，如何准确预测专家负载的方法并不在此仓库的讨论范围内。一种常见的方法是使用历史统计信息的滑动平均值来预测负载。

算法说明

该负载均衡算法包含两种策略，分别适用于不同的场景。

层次化负载均衡

当服务器节点的数量能够整除专家分组的数量时，我们采用层次化负载均衡策略，以充分利用分组受限专家路由机制。首先，我们将专家分组均匀地分配到各个节点上，确保各节点的负载均衡；接着，在每个节点内部对专家进行复制；最后，再将复制后的专家分配到各个 GPU 上，以实现不同 GPU 之间的负载均衡。层次化负载均衡策略适用于预填充阶段，此时专家并行规模较小。

全局负载均衡

在其他情况下，我们则采用全局负载均衡策略，该策略会忽略专家分组，对所有专家进行全局复制，并将复制后的专家分配到各个 GPU 上。此策略适用于解码阶段，此时专家并行规模较大。

接口与示例

负载均衡器的主要函数为 eplb.rebalance_experts。

以下代码展示了一个两层 MoE 模型的示例，每层包含 12 个专家。我们在每层引入 4 个冗余专家，总共生成 16 个副本，并将其放置在 2 个节点上，每个节点配备 4 个 GPU。

import torch
import eplb

weight = torch.tensor([[ 90, 132,  40,  61, 104, 165,  39,   4,  73,  56, 183,  86],
                       [ 20, 107, 104,  64,  19, 197, 187, 157, 172,  86,  16,  27]])

num_replicas = 16
num_groups = 4
num_nodes = 2
num_gpus = 8

phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)
print(phy2log)

# 输出：
# tensor([[ 5,  6,  5,  7,  8,  4,  3,  4, 10,  9, 10,  2,  0,  1, 11,  1],
#         [ 7, 10,  6,  8,  6, 11,  8,  9,  2,  4,  5,  1,  5,  0,  3,  1]])

输出结果由层次化负载均衡策略生成，表明了如下专家复制与放置方案。

许可证

本代码库采用 MIT 许可证发布。

EPLB 快速上手指南

EPLB (Expert Parallelism Load Balancer) 是专为 DeepSeek-V3 等混合专家模型（MoE）设计的专家并行负载均衡工具。它通过冗余专家策略和启发式打包算法，解决不同专家负载不均导致的 GPU 利用率差异问题，并尽可能将同组专家部署在同一节点以减少跨节点通信。

环境准备

• 系统要求: Linux 操作系统
• 前置依赖:
  • Python 3.8+
  • PyTorch (建议版本与您的模型训练/推理环境一致)
  • NumPy

确保已安装 PyTorch。国内用户推荐使用清华或阿里镜像源加速安装：

pip install torch numpy -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

本项目核心逻辑位于 eplb.py 文件中。您可以直接克隆仓库或下载该文件到项目目录。

方式一：克隆仓库

git clone https://github.com/deepseek-ai/EPLB.git
cd EPLB

方式二：直接使用单文件

如果您只需使用负载均衡算法，可直接将 repo 中的 eplb.py 复制到您的项目根目录，无需额外安装步骤。

# 假设您已下载 eplb.py 到当前目录
ls eplb.py

基本使用

EPLB 的核心接口是 eplb.rebalance_experts。该函数接收专家负载估计值，输出平衡后的专家复制与放置方案。

使用示例

以下示例展示了一个两层 MoE 模型的配置：每层包含 12 个专家，引入 4 个冗余专家（共 16 个副本），部署在 2 个节点（共 8 张 GPU）上。

import torch
import eplb

# 模拟专家负载权重 (2 层 x 12 专家)
# 实际场景中，这些数据通常来自历史统计的移动平均值
weight = torch.tensor([[ 90, 132,  40,  61, 104, 165,  39,   4,  73,  56, 183,  86],
                       [ 20, 107, 104,  64,  19, 197, 187, 157, 172,  86,  16,  27]])

# 配置参数
num_replicas = 16  # 每层总副本数 (原始专家 + 冗余专家)
num_groups = 4     # 专家组数量 (用于组限制路由)
num_nodes = 2      # 服务器节点数
num_gpus = 8       # 总 GPU 数

# 执行重平衡
# phy2log: 物理 GPU 索引 -> 逻辑专家索引映射
# log2phy: 逻辑专家索引 -> 物理 GPU 索引映射
# logcnt: 每个逻辑专家的副本计数
phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)

print("物理 GPU 到逻辑专家的映射计划:")
print(phy2log)

输出说明： phy2log 张量展示了最终的部署方案。每一行代表一个模型层，列对应物理 GPU 索引（0-7），值为该 GPU 上运行的逻辑专家 ID。算法会自动根据负载情况选择分层负载均衡（适用于 Prefill 阶段，节点数能整除组数）或全局负载均衡（适用于 Decoding 阶段）策略。