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

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

为什么我们需要 EPLB？#

在使用专家并行 (EP) 时，不同的专家被分配到不同的 NPU 上。鉴于不同专家的负载可能会根据当前工作负载而变化，保持各 NPU 之间的负载均衡至关重要。我们采用冗余专家策略，通过复制高负载专家，然后启发式地将这些副本打包到 NPU 上，以确保它们之间的负载均衡。此外，得益于 MoE 模型中使用的组限制专家路由，我们还尝试尽可能将同一组的专家放置在同一节点上，以减少节点间的数据流量。

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

eplb

如何使用 EPLB？#

有关详细信息，请参阅用户指南的 EPLB 章节：如何使用 EPLB

工作原理是什么？#

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. Adaptor Module
Handles registration and adaptation for different MoE model types

  • abstract_adaptor.py
    Abstract base class defining unified registration interfaces for EPLB adapters

  • vllm_adaptor.py
    Implementation supporting Qwen3-MoE and DeepSeek models, standardizing parameter handling for policy algorithms

2. Core Module
Implements core algorithms, updates, and asynchronous processing

  • Policy Submodule
    Load balancing algorithms with factory pattern instantiation

    • policy_abstract.py
      Abstract class for load balancing strategy interfaces

    • policy_dynamic_ep.py
      Default implementation of open-source EPLB paper algorithm

    • policy_dynamic_ep_v2.py
      Enhanced version optimizing expert swaps for low-bandwidth devices (e.g., A2)

    • policy_flashlb.py
      Threshold-based adjustment reducing operational costs through layer-wise fluctuation detection

    • policy_factory.py
      Strategy factory for automatic algorithm instantiation

  • eplb_device_transfer_loader.py
    Manages expert table/weight transmission and updates

  • eplb_utils.py
    Utilities for expert table initialization and mapping

  • eplb_worker.py
    Asynchronous algorithm orchestration and result processing

3. 系统组件

  • eplb_updator.py
    Central coordinator for load balancing during inference workflows

  • utils.py
    General utilities for EPLB interface registration

关键优化点：

  1. 保持原有结构，同时提高技术清晰度

  2. 标准化术语

  3. 通过简洁的描述增强算法区分度

  4. 通过层次化呈现改进作用域界定

  5. 在优化可读性的同时保留文件/类的关系

默认算法#

层级负载均衡#

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

全局负载均衡#

在其他情况下，我们使用全局负载均衡策略，该策略忽略专家组并在全局范围内复制专家，并将复制的专家打包到单个 NPU 上。此策略可以在解码 (decoding) 阶段采用，配合较大的专家并行规模。

添加新的 EPLB 策略#

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

  1. Inherit the EplbPolicy abstract class of policy_abstract.py and override the rebalance_experts interface, ensuring consistent input parameters current_expert_table, expert_workload and return types newplacement. For example:

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

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

添加新的 MoE 模型#

模型集成实现指南

  1. 适配器文件修改

    • 继承或修改 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

  2. MoE 特性集成

    • 使用 MoE 特定的方法扩展 vllm_ascend/eplb/utils.py

    • 实现专家路由或权重管理所需的功能

  3. 注册逻辑更新

    • 在 model_register 函数中添加补丁 (patch) 逻辑

    • 保持与现有模型类型的向后兼容性

  4. 验证与测试

    • 验证跨层的参数一致性

    • 测试专家表的跨设备通信

    • 与基准实现（如 Qwen3-MoE）进行性能对比测试

关键实现说明：

  • 保留抽象类中现有的接口契约

  • 使用装饰器进行非侵入式补丁集成

  • 利用 eplb_utils.py 进行共享的专家映射操作

DFX (可维护性设计)#

参数校验#

整数参数#

所有整数输入参数必须明确指定其最大值和最小值，并接受有效值校验。例如，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}")

文件路径#

必须检查 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}"
            )

函数规范#

初始化函数#

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

通用函数#

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

一致性#

专家映射表 (Expert Map)#

专家映射表在初始化和更新期间必须全局唯一。在多节点初始化场景下，应使用分布式通信来校验各 rank 间专家映射表的一致性。如果出现不一致，应通知用户哪些 rank 的映射表存在冲突。在更新过程中，如果仅更改了几层或某个 rank 的专家表，更新后的专家表必须与 EPLB 的上下文同步，以确保全局一致性。

专家权重#

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

使用限制#

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

目录