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

Expert Parallelism Load Balancer (EPLB) — Expert Parallelism Load Balancer (EPLB)#

为什么我们需要 EPLB？ — Why We Need EPLB?#

在使用专家并行（EP）时，不同的专家被分配到不同的 NPU。考虑到不同专家的负载可能因当前工作量而异，因此保持不同 NPU 之间的负载平衡至关重要。我们采用冗余专家策略，通过复制负载较重的专家来实现。然后，我们在 NPU 上启发式地打包这些复制的专家，以确保它们之间的负载平衡。此外，由于 MoE 模型中使用了组限制专家路由，我们还会尽可能地尝试将同一组的专家放置在同一节点上，以减少节点间的数据流量。

为了便于复现和部署，Vllm Ascend 支持在 vllm_ascend/eplb/core/policy 中部署 EP 负载均衡算法。该算法根据估计的专家负载计算出均衡的专家复制和放置计划。请注意，预测专家负载的确切方法超出了本仓库的范围。一种常用方法是使用历史统计数据的移动平均。

eplb

如何使用 EPLB？ — How to Use EPLB?#

有关详细信息，请参阅用户指南中的 EPLB 部分： How to Use EPLB

它是如何工作的？ — How It Works?#

EPLB 模块架构

vllm_ascend
├── eplb
│   ├── adaptor
│   │   ├── abstract_adaptor.py
│   │   ├── vllm_adaptor.py
│   ├── core
│   │   ├── policy
│   │   │   ├── policy_abstract.py
│   │   │   ├── policy_dynamic_ep.py
│   │   │   ├── policy_dynamic_ep_v2.py
│   │   │   ├── policy_factory.py
│   │   │   ├── policy_flashlb.py
│   │   ├── eplb_device_transfer_loader.py
│   │   ├── eplb_utils.py
│   │   ├── eplb_worker.py
│   ├── eplb_updator.py
│   ├── utils.py
└───────────

1. 适配器模块
处理不同 MoE 模型类型的注册和适配

abstract_adaptor.py
抽象基类，定义了 EPLB 适配器的统一注册接口
vllm_adaptor.py
支持 Qwen3-MoE 和 DeepSeek 模型的实现，为策略算法标准化参数处理

2. 核心模块
实现核心算法、更新和异步处理

策略子模块
负载均衡算法，采用工厂模式实例化
- policy_abstract.py
  负载均衡策略接口的抽象类
- policy_dynamic_ep.py
  开源 EPLB 论文算法的默认实现
- policy_dynamic_ep_v2.py
  增强版，优化了低带宽设备（例如 A2）的专家交换
- policy_flashlb.py
  基于阈值调整，通过层级波动检测降低运营成本
- policy_factory.py
  用于算法自动实例化的策略工厂
eplb_device_transfer_loader.py
管理专家表/权重的传输和更新
eplb_utils.py
用于专家表初始化和映射的实用工具
eplb_worker.py
异步算法编排和结果处理

3. 系统组件

eplb_updator.py
推理工作流中负载均衡的中央协调器
utils.py
用于 EPLB 接口注册的通用实用工具

关键优化

在保持原始结构的同时，提高了技术清晰度
标准化术语
通过简洁的描述增强了算法区分度
通过分层呈现改进了范围界定
在优化可读性的同时保留了文件/类关系

默认算法 — Default Algorithm#

分层负载均衡 — Hierarchical Load Balancing#

当服务器节点数量能整除专家组数量时，我们使用分层负载均衡策略来利用组限制专家路由。我们首先将专家组均匀地打包到节点上，确保不同节点之间的负载均衡。然后，我们在每个节点内复制专家。最后，我们将复制的专家打包到各个 NPU 上，以确保它们之间的负载均衡。分层负载均衡策略可以在预填充阶段使用，并具有较小的专家并行大小。

全局负载均衡 — Global Load Balancing#

在其他情况下，我们使用全局负载均衡策略，该策略会全局复制专家，而不考虑专家组，并将复制的专家打包到各个 NPU 上。此策略可以在解码阶段采用，并具有较大的专家并行大小。

添加新的 EPLB 策略 — Add a New EPLB Policy#

如果您想向 vllm_ascend 添加新的 eplb 策略，必须遵循以下步骤：

继承 policy_abstract.py 的 EplbPolicy 抽象类，并重写 rebalance_experts 接口，确保输入参数 current_expert_table、expert_workload 和返回类型 newplacement 保持一致。例如：

class RandomLoadBalance(EplbPolicy):

    def __init__(self, config: DynamicConfig):
        super().__init__(config)

    def rebalance_experts(self, current_expert_table, expert_workload):
        new_table = copy.deepcopy(current_expert_table)
        num_layers = len(current_expert_table)

        for i in range(num_layers):
            # randomly choose two card
            # indices = random.sample(range(num_card), 2)
            indices = [3, 1]

            # swap redundant experts
            expert_id_to_exchange = new_table[i][indices[0]][-1].clone()
            new_table[i][indices[0]][-1] = new_table[i][indices[1]][-1]
            new_table[i][indices[1]][-1] = expert_id_to_exchange

        return 1, [-i for i in range(num_layers)], new_table

要添加新的 EPLB 算法，请将策略类型及其对应的实现类包含在 policy_factory.py 的 PolicyFactory 中。

添加新的 MoE 模型 — Add a New MoE Model#

模型集成实现指南

适配器文件修改
- 继承或修改 vllm_ascend/eplb/adaptor/vllm_adaptor.py
- 添加关键参数的处理逻辑
  - num_dense_layers
  - global_expert_num
  - num_roe_layers
- 确保在 model_register 函数中进行参数同步。
  
  例如
  
  修改 vllm_adaptor.py 的 __init__ 以添加新的 MoE 模型 EPLB 参数
```
   if self.model.config.model_type == "qwen3_moe":
    self.num_dense_layers = 0
    self.global_expert_num = self.model.config.num_experts
```
  修改 vllm_adaptor.py 的 model_register 以注册新 MoE 模型的 EPLB 参数
```
    if config.model_type == "qwen3_moe":
        model.num_moe_layers = config.num_hidden_layers
```
MoE 特性集成
- 使用 MoE 特定方法扩展 vllm_ascend/eplb/utils.py
- 实现专家路由或权重管理所需的函数
注册逻辑更新
- 在 model_register 函数中添加补丁逻辑
- 保持与现有模型类型的向后兼容性
验证与测试
- 验证跨层参数的一致性
- 测试专家表的跨设备通信
- 与基线实现（例如 Qwen3-MoE）进行基准测试

关键实现注意事项

在抽象类中保留现有的接口契约
使用装饰器进行非侵入性补丁集成
利用 eplb_utils.py 进行共享专家映射操作

DFX#

参数验证 — Parameter Validation#

整数参数 — Integer Parameters#

所有整数输入参数必须明确指定其最大值和最小值，并接受有效值验证。例如，num_iterations_eplb_update 必须大于 0。

    @staticmethod
    def check_iterations(iterations):
        if not isinstance(iterations, int):
            raise TypeError(f"The {iterations} is not int.")
        if iterations <= 0:
            raise ValueError(
                f"The {iterations} can not less than or equal to 0.")
        if iterations > sys.maxsize:
            raise ValueError(
                f"The {iterations} can not large than {sys.maxsize}")

文件路径 — File Path#

必须检查 EPLB 的文件路径的合法性，例如路径是否有效以及是否具有适当的读写权限。例如：

    @staticmethod
    def check_expert_map_path(expert_map):
        if expert_map is None:
            return
        if not isinstance(expert_map, str):
            raise TypeError("The expert_map is not str.")
        if not expert_map.strip():
            raise ValueError("The expert_map is not empty.")
        _, ext = os.path.splitext(expert_map)
        if ext.lower() != ".json":
            raise TypeError("The expert_map is not json.")
        if not os.path.exists(expert_map):
            raise ValueError("The expert_map is not exist.")
        try:
            with open(expert_map, "w", encoding='utf-8') as f:
                f.read()
        except Exception as e:
            raise IOError(
                f"Fail read expert info from {expert_map}, please check the reading permission of {expert_map} : {e}"
            )

函数规范 — Function Specifications#

初始化函数 — Initialization Function#

所有 EPLB 参数必须在初始化时默认初始化，并指定参数类型和默认值以供正确处理。

通用函数 — General Functions#

所有方法参数都必须指定参数类型和默认值，函数必须为默认参数包含默认返回值处理。建议使用 try-except 块处理函数体，指定捕获的异常类型和失败处理（例如，记录异常或返回失败状态）。

一致性 — Consistency#

专家映射 — Expert Map#

专家映射在初始化和更新期间必须是全局唯一的。在初始化期间的多节点场景下，应使用分布式通信来验证每个 rank 的专家映射的一致性。如果不一致，应通知用户哪些 rank 存在不一致的映射。在更新过程中，如果只有少数层或某个 rank 的专家表已更改，则必须将更新后的专家表与 EPLB 的上下文同步，以确保全局一致性。

专家权重 — Expert Weight#

更新专家权重时，请确保已释放专家权重的内存分配，或者专家（指旧版本）不再使用。

限制#

在使用 EPLB 之前，启动脚本并添加 export DYNAMIC_EPLB="true"。在执行负载数据收集（或性能数据收集）之前，启动脚本并添加 export EXPERT_MAP_RECORD="true"。