摘要：EPLB（Expert Parallelism Load Balancer）是一种专为 Mixture-of-Experts (MoE) 模型设计的负载均衡算法，旨在优化专家的复制与分布策略，以最大化 GPU 计算效率并减少跨节点通信。其核心创新点包括 “冗余专家” 复制机制，通过动态增加高负载专家的副本来均衡计算资源，并采用 “分层负载均衡”（Hierarchical Load Balancing）策略，使专家分组优先匹配计算节点，减少通信开销。此外，在无法整除分配的情况下，EPLB 还提供 “全局负载均衡”（Global Load Balancing）方案，以确保所有 GPU 负载均衡。
本文深入解析 EPLB 的关键函数及其实现方式，包括 balanced_packing（贪心式专家分配）、replicate_experts（智能专家复制）、rebalance_experts（负载均衡主函数）等。同时，我们探讨了 EPLB 在 MoE 训练和推理中的应用，并提出了一些优化建议，如提高算法效率、改进非均匀硬件适配、实现在线动态负载均衡等。总体而言，EPLB 通过“复制+打包”的创新方法，成功解决了 MoE 训练中的负载不均衡问题，为大规模 AI 模型的高效部署提供了强有力的支撑。

1. 功能解析

（1）balanced_packing 函数：这个函数用于把一定数量的“物品”按照权重分配到固定数量的“箱子”（pack）中，使每个箱子装入的物品数量相同且总重量尽可能平衡。它的参数 weight 是形状 [X, n] 的张量，每行代表一层（或一次试验）里 n 个物品的权重；num_packs 是箱子数。函数返回两个张量：pack_index（每个物品被放入哪一个箱子的索引）和 rank_in_pack（物品在箱子中的位置）。核心算法是贪心法：先对每行的物品按权重降序排序，然后依次将最重的物品放入当前总重量最轻且仍未满的箱子中。每个箱子最多能放 n/num_packs 个物品（题目保证 n 能被 num_packs 整除）。这样循环分配，直到每个物品都有箱子。举例来说，如果某层有4个专家组负载 [9,7,5,3]，需要分到2个节点（每节点2个组）：算法会先将权重9分配给节点0（此时节点0总负载9），再将7分配给节点1（节点1负载7），接着下一个重量5应分给当前负载较轻的节点1（节点0=9，节点1=7），分配后节点1负载变为12；最后权重3给节点0（节点0负载变为12）。结果是每个节点分到2个组且总负载均为12，达到了负载平衡。balanced_packing 的实现中也考虑了特殊情况：如果每个箱子只需装1个物品（如 groups_per_pack == 1），则直接按顺序分配即可，无需排序。

（2）replicate_experts 函数：这个函数根据各专家的负载权重，将“重负载”的专家复制（冗余）出额外的副本，从而使总的物理专家数量达到 num_phy，并尽量使每个专家副本承担的负载均衡。参数 weight 是形状 [X, num_log] 的张量，表示 X 层中每个逻辑专家（原始专家）的负载；num_phy 是希望达到的物理专家总数（即原始专家加上冗余副本的总数）。返回结果包括：phy2log（形状 [X, num_phy]，表示每个物理专家对应的逻辑专家ID）、rank（每个物理专家是其对应逻辑专家的第几个副本，称为副本序号）以及 logcnt（形状 [X, num_log]，每个逻辑专家被复制的总次数）。

实现思路是迭代复制最繁忙的专家：最开始假定每个逻辑专家各有一个物理实例（副本），所以 logcnt 初始全为1。然后在需要添加的 num_redundant = num_phy - num_log 个副本中，每次选择当前“平均负载”最高的专家来复制。具体做法是：在每次迭代中计算 weight / logcnt（即每个专家每个副本的平均负载），找到值最大的专家索引，将该专家再复制一次：为它新增一个物理ID，并记录这个物理专家对应的逻辑专家编号（填入 phy2log）以及它是该专家的第几个副本（填入对应的 rank，这个值等于复制前该专家已有的副本数）。然后将对应专家的计数 logcnt 加1，并进入下一轮迭代。这样重复，直到创建了所需数量的冗余副本。

示例：假设某层有3个专家的负载分别为 [50, 30, 20]，希望扩展到5个物理专家（即复制2个额外的专家副本）。初始时每个专家各有1个副本（logcnt = [1,1,1]），物理专家ID 0、1、2 分别对应逻辑专家0、1、2（phy2log 初始为 [0,1,2,_,_]，下划线表示待填充）。现在需要添加第3个物理专家：比较各专家的平均负载 50/1, 30/1, 20/1，第0号专家负载最大（50），于是复制专家0，赋予新物理ID 3，并令 phy2log[3]=0。这个副本是专家0的第2个实例，所以 rank[3]=1（第一个副本的 rank=0），更新专家0的计数为2（logcnt[0]=2）。接着添加第4个物理专家：此时平均负载变为 专家0:50/2=25, 专家1:30/1=30, 专家2:20/1=20，最大的是专家1，于是复制专家1（物理ID 4 对应专家1，phy2log[4]=1），它是专家1的第2个副本（rank[4]=1），更新 logcnt[1]=2。最终得到的映射：phy2log = [0,1,2,0,1]，表示物理专家0和3都对应逻辑专家0，物理专家1和4对应专家1，物理专家2对应专家2；相应的 rank = [0,0,0,1,1]，其中物理专家3是专家0的副本1号，物理专家4是专家1的副本1号；logcnt = [2,2,1]，表示专家0和1各有2个物理副本，专家2仍为1个。这种方法通过为负载最大的专家添加额外副本，降低了单个副本需要处理的平均负载，从而使最大负载得到平衡。

（3）rebalance_experts 函数：这是 EPLB 的主接口函数，执行专家负载均衡策略。它根据专家负载 weight（形状 [layers, num_logical_experts]，表示每层每个逻辑专家的负载），以及集群的配置参数（num_replicas 目标物理专家总数，num_groups 专家分组数，num_nodes 服务器节点数，num_gpus GPU总数）来计算专家复制和放置方案。返回三个张量：physical_to_logical_map (phy2log，每层每个物理专家对应的逻辑专家编号)、logical_to_physical_map (log2phy，每层每个逻辑专家的各个物理副本索引列表) 和 expert_count (logcnt，每层每个逻辑专家被复制的数量)。

rebalance_experts 内部根据情况采用分层（分级）的负载均衡或全局负载均衡两种策略：

• 如果专家组数 num_groups 能被节点数 num_nodes 整除（即每个节点可分配若干完整的专家组），则采用**“分层负载均衡”**策略（Hierarchical Load Balancing）。这种情况下，函数会调用内部的 rebalance_experts_hierarchical 来执行三步操作：

  1. 组到节点的负载平衡：将专家按组打包并分配到各节点，使用的正是前述 balanced_packing 算法。在分组路由的MoE中，通常多个专家组成一个组，推理时每个 token 只路由到某一组内的专家。为了减少节点间通信，算法倾向于将同一组的专家尽可能放在同一节点。具体实现中，首先计算每组专家的总负载（把 weight 按每组大小汇总），然后用 balanced_packing(…, num_nodes) 将这些组分配给不同节点，使每个节点承担的总组负载尽量均衡，且每个节点分配相同数量的组。这样得到每层“组->节点”的映射（即 group_pack_index 等）。这一步确保各节点总的专家组负载相近。

  2. 节点内专家复制（冗余）：接下来，在每个节点内部独立地进行专家复制负载均衡。也就是对每个节点上属于它的那些专家，将它们的负载提取出来，然后调用 replicate_experts（针对该节点的专家集合）来决定在该节点内如何复制重负载专家，以及这些副本在节点内的映射。这一步产生每个节点内的物理专家映射 phy2mlog（物理专家对应节点内逻辑专家）的列表，以及节点内每个专家的副本计数 mlogcnt。因为各节点之间此时相互独立，这一步确保每个节点自身的专家负载得到均衡（不会因为某节点有过多高负载专家却无法复制而失衡）。

  3. 节点内副本分配到GPU：最后，将各节点内产生的物理专家进一步分配到该节点的各个GPU上，使节点内的每块GPU负载平衡。每个节点内我们已经有一定数量的物理专家（包括复制的），且每个节点有 num_gpus/num_nodes 块GPU。算法再次利用 balanced_packing：根据每个物理专家负责的平均负载（专家总负载除以副本数得到每个副本的“权重”），将这些物理专家分配到固定数量的GPU（每GPU分到等量的物理专家副本），使各GPU承担的总负载尽量相等。这会返回物理专家到GPU的打包索引（pack index）。据此，算法可以确定每个物理专家的最终“物理位置”（哪一个GPU），并生成最终的 physical_to_logical_map 映射，以及对应的 logical_to_physical_map 和 expert_count。

• 如果专家组数无法与节点数整除（即无法完整地将组平均分配给节点），则采用**“全局负载均衡”**策略（Global Load Balancing）。这种情况下，不考虑组的划分，直接在全局范围对专家进行复制和平衡。实现上，rebalance_experts 会调用前述的 replicate_experts(weight, num_replicas) 获取在全局范围内应该复制哪些专家、复制多少次。这样得到的 phy2log 和 logcnt 直接体现哪些专家有多个副本。由于全局策略下每个专家副本可以放在任意GPU上，为了输出方便，代码还构造了 logical_to_physical_map：它创建一个大小为 [layers, num_logical_experts, max_copies] 的张量，将每个逻辑专家的副本物理ID填进去（没有用到的位置填 -1）。在全局策略中，由于输入参数要求 num_replicas 是 num_gpus 的整数倍，因此可以简单地将这些物理专家平均分配给所有GPU（例如按顺序每台GPU拿相等数量的副本）来保证负载均衡——这一点由调用方根据返回的映射去实现。总的来说，全局策略更加直接，忽略了专家组的概念，当分组与节点拓扑不匹配时采用该策略以保证所有专家全局范围内负载均衡。

（4）代码示例：官方提供了一个两层MoE模型的示例，每层12个专家，冗余4个专家副本，总共16个物理专家放置在2个节点（每节点4块GPU）上。调用：

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[0])  # 打印第0层的 physical_to_logical 映射

输出（第0层的映射）示例：tensor([5, 6, 5, 7, 8, 4, 3, 4, 10, 9, 10, 2, 0, 1, 11, 1])。这表示第0层共有16个物理专家（索引0到15），其中物理ID为0和2的都是逻辑专家5（Expert 5被复制了两次）、ID为1的是专家6、ID3是专家7，等等。可以看到原本负载很高的专家5、专家10、专家1等在映射中出现了不止一次，说明它们被复制成多个实例以分摊负载。而每个GPU分到了2个物理专家，确保了8块GPU之间负载均衡。在第1层，类似地高负载的专家（例如专家5、6、8等）也被复制。最终得到的 logcnt 会是每层每个专家的复制数（例如第0层 Expert5的计数为2，Expert10为2，Expert1为2，等等）。logical_to_physical_map 则能列出每个专家对应的物理ID列表。例如对于第0层 Expert5，可能对应物理ID [0, 2]（表示Expert5对应物理0号和2号副本）。

综上，rebalance_experts 利用以上子函数，在给定专家负载统计的情况下，自动计算一个专家副本的增加和放置方案，使GPU之间、节点之间的负载都尽量平衡，同时尽可能利用专家分组信息减少跨节点通信。

2. 创新点

冗余专家的负载均衡策略：EPLB 的核心创新在于引入“冗余专家”策略和分层负载均衡算法。在传统的 Mixture-of-Experts (MoE) 系统中，负载均衡通常通过调整路由算法（如限制每个专家接收的 token 数）或者损失函数中的均衡项来实现。而 EPLB 则另辟蹊径，通过复制高负载专家的方式，从系统架构层面增加容量来分担负载。也就是说，如果某些专家工作负载过高，就动态地增加它的副本，让不同的 GPU 可以并行处理该专家的不同一部分工作。这种方法避免了简单地重新分配已有专家所带来的局限，在保证模型能力的同时，有效缓解个别专家的过载问题。可以认为，这是一种“以空间换时间”的策略，利用额外的计算实例来提升并行度，从而达到负载平衡。

分层（分级）负载均衡 vs. 全局负载均衡：EPLB 提供了两种策略，根据集群拓扑和专家分组情况选择使用：

• 分层负载均衡（Hierarchical LB）：这个策略的独特之处在于充分利用了专家分组与节点拓扑的对应关系。在一些 MoE 模型（如 DeepSeek-V3）中，引入了“组限路由”（group-limited routing）的概念，即将专家划分为若干组，每个数据样本只会路由到某一组内的专家。EPLB 捕捉到了这一点，优先将同一组的专家尽量放在同一个物理节点上。这样做的好处是：组内通信留在节点内部（比如利用高速的NVLink），不同组之间尽可能减少跨节点的通信开销。EPLB 的层次化策略首先在组粒度上保证节点间均衡，再在节点内部复制专家、均衡GPU，这种“两级平衡”的过程充分考虑了通信局部性和负载均衡的结合。这种算法确保如果集群配置允许，每个节点处理相对独立的一部分专家组，节点内部各GPU再细化分工，从而达到全局的高效均衡。这一点在以往的负载均衡方案中是少见的，将系统架构（节点/GPU拓扑）与模型结构（专家组划分）结合起来，是 EPLB 的一大特色。

• 全局负载均衡（Global LB）：当专家的分组与节点数不契合时，EPLB 切换到全局策略。其独特之处在于以全局视角考虑所有专家的负载，忽略拓扑限制，直接通过复制平衡每个专家的负载。这虽然没有考虑通信局部性，但在大型推理阶段（如解码阶段专家并行规模很大）更实用。全局策略可以被视为对任意情况下的通用负载均衡：它确保在总体上，每个 GPU 分到的专家实例数相等，每个专家的工作量通过复制被均摊。这种方法简单有效，兼容性强。当分组路由无法很好利用时，全局复制策略依然能够提供可靠的负载均衡方案。相较于层次化策略，全局策略更直接和通用。

小结：EPLB 的创新点在于**“复制+打包”的思路：一方面通过冗余复制专家解决专家负载不均的问题，另一方面通过贪心打包算法在不同层级（节点级、GPU级）上分配专家实例，实现负载均衡。它综合考虑了模型层（专家负载）和系统层（多节点多GPU拓扑）的信息。例如，将组内专家固定在单节点就是在降低通信开销的同时，还通过两步打包确保各节点、各GPU负载相近。这种多层次的均衡策略使 EPLB 相比简单均衡方法更加灵活高效**。在DeepSeek-V3等实际应用中，这一方法体现出在不同阶段（预填充、小规模并行；解码、大规模并行）动态选择策略的优势，保证了模型推理速度和硬件利用率。

3. 优化方案

尽管 EPLB 提供了有效的负载均衡方案，但仍有一些值得探索的改进方向：

（1）算法效率优化：当前实现中的某些步骤是通过 Python 循环完成的，对于大量专家时可能成为瓶颈。例如，balanced_packing 函数对每层的专家组排序并逐个分配，如果专家组数目很大，排序是 $O(n\log n)$，分配则是 $O(n)$，总体对每层是线性对数级。replicate_experts 函数在需要添加大量冗余副本时，会在循环中每次计算一次全体专家的 weight/logcnt 最大值，这样每添加一个副本都是一次 O(n) 扫描，循环执行 $num\_redundant$ 次，最坏情况时间复杂度接近 $O(n\times num\_redundant)$。在极端情况下（专家很多且需要复制很多次）效率可能较低。可以考虑的优化包括：

• 利用优先队列（heap）或其它数据结构加速寻找当前最大平均负载的专家。比如将所有专家按当前 weight/logcnt 保存在一个最大堆中，每次弹出最大者进行复制，并更新其值后再插入堆。这样每次复制操作复杂度可降为 $O(\log n)$，总体复杂度约 $O(n+num\_redundant\log n)$。
• 将关键步骤移至底层实现或GPU并行化。例如，balanced_packing 的贪心分配可以尝试用向量化的方式处理多层的数据，或者使用CUDA核函数并行处理每层的分配，从而减少Python层的开销。
• 针对排序等静态步骤，可以在负载变化不大时重用以前的排序结果或部分更新，而不每次从头计算。例如若负载变化幅度有限，可考虑增量式地调整分配。

（2）代码结构和可读性优化：目前的实现虽然在功能上清晰，但一些步骤较为繁琐，理解门槛高。例如层次化策略中大量使用张量变换（unflatten、gather、scatter、flatten等）来在不同映射之间转换。为了提高可维护性，可以：

• 增加更多的注释和文档说明，明确每个中间张量的物理含义。如 phy2mlog、pphy2log 等变量，可以在代码中注释说明它们代表的映射关系，方便他人阅读和修改。
• 尝试将部分逻辑拆分成小函数，或用更具语义的变量名，使过程更直观。例如将“节点内复制”或“GPU打包”各部分分成独立函数，名字直接反映其功能。
• 在接口层面，确保无论是层次化还是全局策略都返回一致的输出格式。目前实现中，rebalance_experts 对全局策略直接构造了 log2phy，但对层次化策略返回时并未显式组装 log2phy（虽然通过 phy2log 和 rank 也可以重构）。为了统一和方便，或可在层次化分支也生成完整的 logical_to_physical_map。这类改进不会影响算法结果，但可以避免使用者因输出格式差异产生困扰。

（3）负载均衡策略增强：

• 更智能的复制决策：当前的冗余复制策略是贪心的局部最优，可能不是全局最优。例如，在某些负载分布下，贪心算法选择的一系列复制顺序可能导致接近最优但略有差距的结果。未来可以研究使用线性规划或动态规划来近似求解“如何分配副本使最大负载最小”的问题，或者引入对专家负载更全面的考量（如结合延迟、通信成本等）。这可能提升全局负载平衡的效果。
• 非均匀硬件情况下的平衡：EPLB 假设每个GPU计算能力相同，每台服务器的GPU数量也相等。但在实际部署中，有时硬件性能不均衡（例如不同节点GPU数不同，或GPU性能不同）。策略可以扩展以接受权重或容量参数，为每个节点/GPU设定不同的“容量权重”，按此进行非均匀的负载平衡。比如修改打包算法，让每个箱子的容量权值不同，使得强性能设备承担更多副本，而总体仍平衡。
• 动态负载感知和调整：当前算法基于预估的负载（例如历史平均）。在实际运行中，负载可能实时变化。可以考虑将 EPLB 融入一个动态调整系统：定期监控各专家的实际负载，如果失衡再次出现，就触发 rebalance_experts 重新计算新的复制和放置方案，实现在线负载均衡。为降低抖动和开销，可以在复制和回收副本之间做平滑处理，例如仅当某专家持续过载一段时间才增加副本，空闲时再逐步回收。此外，也可引入滞后或阈值，避免频繁重分配。
• 更广泛的策略兼容：对于 num_groups 不是 num_nodes 整数倍的情况，目前简单地使用全局策略。或许可以探索混合策略：例如大部分专家组整除分配，用层次化方法放置，余下少量“无法整除的”专家组用全局方法处理，尽可能仍将组内专家放在同节点。这样在不完全匹配时也能部分利用组局部性。虽增加复杂度，但可能在某些配置下提供更优的通信和负载权衡。

（4）其他实现细节优化：

• 在内存使用上，注意大张量的创建和复制。例如全局策略中构造 log2phy 张量时，尺寸为 [layers, num_experts, max_copies]，其中很多元素可能填充为 -1 没有实际用处。如果专家副本数差异很大，这个张量会比较稀疏。可以考虑用稀疏表示或者仅记录实际索引列表的方式来替代完整的三维张量，减少内存占用。
• 考虑随机性或扰动引入：当前算法是基于确定的负载排序和选择，可能会在某些情况下形成固定模式（比如总是某些特定专家被复制）。在某些训练或推理场景下，适当的随机扰动（如在权重相近时随机选一个复制）可能避免潜在的性能瓶颈或过拟合某一方案。这方面虽不是主要矛盾，但在扩展研究中可以尝试。

总而言之，EPLB 已经提供了一个有效的专家并行负载均衡框架。通过上述优化方案，有望进一步提升其性能和适应性：既能处理更大规模的专家和更复杂的硬件拓扑，又能以更高效的实现降低开销，并在更多场景下保持专家负载和通信的平衡。这些改进将使其在Mixture-of-Experts模型的部署中更加实用和健壮。