vllm源码剖析15-vLLM 分布式推理-EPLB负载均衡 文章目录一 EPLB 模块概述1.1 背景知识1.2 基本概念1.3 核心思想二 专家并行 EPLB 参数和服务实例三 vLLM EPLB 设计方案四 vLLM EPLB 流程解析参考资料一 EPLB 模块概述1.1 背景知识我们知道混合专家模型Mixture-of-Experts, MoE会通过路由机制将不同的输入 token 分配给不同的专家Expert处理。每个专家通常可以理解为一组类似 MLP 的前馈网络结构。对于 top-k 路由的 MoE 层一个 token 可能被分配给多个专家最终输出通常会按照路由权重对这些专家结果进行加权汇聚而不是简单做无权重的逐元素相加。在部署 MoE 模型时通常会结合专家并行Expert Parallelism, EP和数据并行等策略来加速推理。专家并行简单来说就是把 MoE 层中的专家分布到不同 GPU 或 rank 上让每个 rank 主要负责本地专家的计算。在 vLLM 中EPLB 也是建立在启用专家并行的前提下工作的当开启 EPLB 时系统会根据专家负载统计调整逻辑专家到物理专家副本的映射并在需要时重排专家权重。专家并行Expert Parallelism, EP 的优势在于内存效率 每个 GPU 通常只需要存储分配到本地的专家权重而不是复制所有专家权重。计算效率多个 GPU 可以同时处理不同专家上的 token 计算这一点和 TP 一样都是利用多卡并行提升吞吐。扩展性可以通过增加 GPU 数量承载更多专家或更大的 MoE 模型。但在 MoE 模型中不同专家接收的任务量tokens 数量可能存在显著差异这会导致专家计算的负载不均衡。具体表现为热门专家Hot Experts处理了大量 tokens所在 GPU 或 rank 的计算负载过高容易成为性能瓶颈。冷门专家Cold Experts处理的 tokens 较少所在 GPU 或 rank 的计算资源可能处于等待或低利用率状态。出现热门专家和冷门专家是因为 MoE 模型会通过路由器为每个 token 选择少量专家。不同 token 的路由结果并不完全相同因此各个专家接收的 token 数量可能存在差异。某些专家接收的 token 较多成为热门专家另一些专家接收的 token 较少成为冷门专家。实际输入本身也不均匀。某个批次可能集中包含相似主题、语言或文本模式于是相关专家会接收到更多 token成为热门专家不常被选中的专家则成为冷门专家。当专家分布在不同 GPU 或 rank 上时热门专家会带来更高的计算负载而冷门专家可能提前完成任务并等待。后续流程通常仍需等待最慢的设备因此热门专家容易成为整体推理性能的瓶颈。EPLBExpert Parallelism Load Balancer就是 vLLM 中用于缓解 MoE 专家并行负载不均衡的机制。它会记录一段窗口内的专家负载并按照配置的步长周期性触发专家重排当配置了冗余专家时还可以通过增加物理专家副本让高负载的逻辑专家被更多物理专家分担从而降低热点专家对整体推理性能的影响。1.2 基本概念在 vLLM 中EPLBExpert Parallel Load Balancing由状态管理、负载均衡策略和专家权重迁移三部分共同完成状态管理EplbState / EplbModelStateEplbState 维护专家负载的滑动窗口并在达到重排间隔时触发重排。EplbModelState 保存单个 MoE 模型的专家负载、映射关系、通信缓冲区和通信器。权重迁移完成后系统会提交新的专家映射。负载均衡策略AbstractEplbPolicy / DefaultEplbPolicyAbstractEplbPolicy 定义统一的重排接口。默认实现 DefaultEplbPolicy 根据逻辑专家的负载和并行配置为负载较高的专家分配更多物理副本并生成新的 physical_to_logical_map。该映射描述每个物理槽位应当存放哪个逻辑专家的权重。策略生成映射后还会尽量保留 GPU 内已有槽位如果某个逻辑专家在重排前后仍然位于同一 GPU就优先让它继续使用原来的物理槽位。这样可以避免仅因槽位顺序变化而复制权重。权重迁移rebalance_execute.py / eplb_communicator.py每个 GPU 上有若干存放专家权重的物理槽位。重排时rebalance_execute.py 根据新旧 physical_to_logical_map 判断各个槽位是否需要更新a. 如果槽位对应的逻辑专家没有变化则保留原有权重。b. 如果目标专家已存在于当前 GPU 的其他槽位则在本地复制权重。c. 如果目标专家不在当前 GPU则通过 eplb_communicator.py 从其他 rank 接收权重。eplb_communicator.py 提供 NCCL、Gloo、PyNCCL 和 NIXL 等通信后端。迁移的数据是 MoE 专家的参数权重例如投影层权重而不是 token 或路由结果。EPLB 算法涉及以下核心概念eplb_communicator.py 提供 NCCL、Gloo、PyNCCL 和 NIXL 等通信后端。迁移的数据是 MoE 专家的参数权重例如投影层权重而不是 token 或路由结果。EPLB 算法涉及以下核心概念专家的数量关系全局物理专家数 逻辑专家数 冗余专家数 每个 GPU 的本地物理专家数 全局物理专家数 / EP ranks 数在 vLLM 的 EPLB 状态中physical_to_logical_map 记录每个物理专家槽位存放哪个逻辑专家的权重形状为 (num_moe_layers, num_physical_experts)logical_to_physical_map 反向记录每个逻辑专家对应哪些物理副本形状为 (num_moe_layers, num_logical_experts, max_replicas)未使用的位置以 -1 表示logical_replica_count 则记录各逻辑专家拥有的副本数量。三者从不同方向描述当前的专家部署关系前者用于负载汇总和权重重排后两者用于获取逻辑专家的可用物理副本。例如假设某个 MoE 层有 4 个逻辑专家 E0、E1、E2、E3以及 6 个物理槽位。当前映射如下physical_to_logical_map [0, 1, 2, 3, 0, 1]这表示物理槽位 P0 和 P4 存放逻辑专家 E0 的权重P1 和 P5 存放 E1 的权重P2、P3 分别存放 E2、E3 的权重。反向映射为logical_to_physical_map [ [0, 4], # E0 对应 P0、P4 [1, 5], # E1 对应 P1、P5 [2, -1], # E2 只有 P2 [3, -1], # E3 只有 P3 ] logical_replica_count [2, 2, 1, 1]中logical_to_physical_map 使用 -1 填充未使用的位置logical_replica_count 则直接记录每个逻辑专家拥有多少个物理副本。三者表达的是同一组专家部署关系physical_to_logical_map 从物理槽位查询逻辑专家另外两个张量从逻辑专家查询可用副本及其数量。负载统计通过 expert_load_pass 和 expert_load_window 完成。expert_load_pass 记录当前 forward pass 中各物理专家处理的 token 数在需要记录负载的步骤结束时数据会写入 expert_load_window。触发重排后系统根据 physical_to_logical_map 将窗口内的物理专家负载汇总为逻辑专家负载据此重新分配副本并生成新的映射。专家权重迁移完成后系统同步更新三类映射张量形成统计负载、计算映射、迁移权重、提交映射的闭环。下面通过一个简化示例串联各个模块。假设某个 MoE 层包含 4 个逻辑专家和 6 个物理槽位分布在两个 GPU 上每个 GPU 存放 3 个槽位。当前映射为GPU 0: [E0, E1, E2] GPU 1: [E3, E0, E1] physical_to_logical_map [0, 1, 2, 3, 0, 1]假设统计结果表明 E0 的负载明显高于其他专家。eplb_state.py 中的 EplbState 会汇总滑动窗口内的负载并调用 DefaultEplbPolicy 计算新的槽位分配。为便于说明假设策略生成的新映射为GPU 0: [E0, E1, E3] GPU 1: [E2, E0, E0] physical_to_logical_map [0, 1, 3, 2, 0, 0]这个结果表示 E0 的物理副本从 2 个增加到 3 个。随后rebalance_execute.py 比较新旧映射并迁移权重GPU 1 的最后一个槽位原本存放 E1现在需要存放 E0而 GPU 1 已经拥有 E0因此可以在本地复制权重GPU 1 的第一个槽位原本存放 E3现在需要存放 E2但 GPU 1 没有 E2因此需要通过 eplb_communicator.py 从其他 rank 接收权重。迁移完成后eplb_state.py 提交新的 physical_to_logical_map并同步更新 logical_to_physical_map 和 logical_replica_count。后续 token 路由将使用新的专家副本布局。vLLM 框架中的 EPLB 模块位于 vllm/distributed/eplb 目录vllm/distributed/eplb/ ├── __init__.py# 模块导出├── eplb_state.py# EPLB 状态管理、负载统计和重排触发├── rebalance_execute.py# 专家权重重排执行├── async_worker.py# 异步 EPLB 传输工作线程├── eplb_communicator.py# 专家权重传输通信实现├── eplb_utils.py# EPLB 辅助工具└── policy/# 负载均衡策略├── __init__.py ├── abstract.py# 抽象策略接口└── default.py# 默认策略实现基于 DeepSeek EPLB 思路1.3 核心思想EPLBExpert Parallelism Load Balancing的核心思想是在专家并行场景中根据一段时间窗口内统计到的专家负载重新计算“物理专家到逻辑专家”的映射并通过冗余专家副本把热点逻辑专家的负载分摊到更多 GPU 上。我们先把术语全部换成“公司”来理解Node节点 一台服务器/一栋楼 GPU 服务器里的显卡/房间 专家 E 放在显卡上的一份特殊权重/员工 专家组 G 一组专家/一个部门 冗余专家 专家的复制品/同岗位的新员工补充①2. 分层策略是什么意思假设 G0 部门被分配给节点 0节点 0 ├── GPU 0G0.E0 └── GPU 1G0.E1 如果 G0.E1 很忙需要复制一个 节点 0 ├── GPU 0G0.E0、G0.E1副本 └── GPU 1G0.E1②全局策略是什么意思全局策略不管“部门属于哪栋楼”专家可以随便放节点 0 ├── GPU 0G0.E0、G1.E0 └── GPU 1G2.E0 节点 1 ├── GPU 2G0.E1、G1.E1 └── GPU 3G2.E1热点专家冗余复制EPLB 不是在运行中无限增加专家数量而是在配置的 num_redundant_experts 物理专家槽位内为高负载逻辑专家分配更多副本。每个副本都可以处理路由到同一逻辑专家的 token从而降低单个物理专家或单张 GPU 的压力。例如假设模型包含 4 个逻辑专家 E0、E1、E2、E3并配置了 2 个冗余槽位因此一共有 6 个物理槽位。初始布局可能是[P0, P1, P2, P3, P4, P5][E0, E1, E2, E3, E0, E1]如果运行期间发现 E2 的负载较高而 E1 的负载较低EPLB 可以重新分配冗余槽位重排前[E0, E1, E2, E3, E0, E1] 重排后[E0, E1, E2, E3, E0, E2]分层负载均衡策略默认策略会优先考虑服务器节点之间的边界。简单来说它不会直接将所有专家打散到所有 GPU 上而是先决定每个专家组应放在哪个节点再在节点内部安排具体 GPU。例如假设集群包含 2 个节点每个节点有 2 张 GPU模型包含 4 个专家组每组有 2 个逻辑专家节点 0GPU 0、GPU 1 节点 1GPU 2、GPU 3 专家组G0、G1、G2、G3策略首先统计每个专家组的总负载并将 4 个 group 均衡地打包到 2 个节点。假设分配结果为节点 0G0、G3 节点 1G1、G2接下来策略分别在每个节点内部完成两件事a. 根据逻辑专家的负载为热点专家分配冗余物理副本。b. 将该节点内的物理专家继续均衡地打包到节点内的 GPU 上。最终可能得到类似布局节点 0 GPU 0G0.E0、G3.E0、G3.E1 GPU 1G0.E1、G0.E1 的冗余副本、G3.E0 的冗余副本 节点 1 GPU 2G1.E0、G1.E1、G2.E0 GPU 3G2.E1、G1.E0 的冗余副本、G2.E1 的冗余副本也就是说一个专家组中的逻辑专家需要增加副本时新副本只会放在该组所属的 Node 内不会跨 Node 扩充。这样做的原因是节点内 GPU 通常可以通过更快的互联通信。将同一专家组及其冗余副本限制在分配到的节点内可以在均衡负载的同时尽量保持局部性减少不必要的跨节点通信。全局负载均衡策略当专家组数量不能被服务器节点数整除时默认策略会退化为全局策略即不再保留专家组到节点的边界约束而是在所有可用 EP rank 上重新平衡物理专家分布。另一个需要注意的情况是如果 GPU 数量不能被节点数整除vLLM 会先将 num_nodes 置为 1因此同样不会使用分层重排策略。例如集群有 2 个 Node但模型包含 3 个专家组Node 0GPU 0、GPU 1 Node 1GPU 2、GPU 3 专家组G0、G1、G2由于 3 个专家组无法均匀分配到 2 个 NodeDefaultEplbPolicy 会退化为全局负载均衡将参数等效地设置为单一分组和单一节点再在全部 4 张 GPU 上重新安排物理专家。此时不再保留专家组到 Node 的边界约束。某个高负载逻辑专家需要增加物理副本时新增副本可以分配到任意可用 GPU因此可能位于任意 Node。二 专家并行 EPLB 参数和服务实例三 vLLM EPLB 设计方案四 vLLM EPLB 流程解析参考资料DeepSeek EPLB 开源库vLLM 官方文档 - Expert ParallelismMoE并行负载均衡EPLB的深度解析与可视化DeepSeek-V3 论文