跳转至

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

为什么需要 EPLB?

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

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

eplb

如何使用 EPLB?

详细信息请参阅用户指南中的 EPLB 部分:如何使用 EPLB

工作原理是什么?

EPLB 模块架构

vllm_ascend
├── eplb
   ├── adaptor
      └── vllm_adaptor.py
   ├── core
      ├── policy
         ├── policy_abstract.py
         ├── policy_default_eplb.py
         ├── policy_factory.py
         ├── policy_flashlb.py
         ├── policy_random.py
         └── policy_swift_balancer.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

  • 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_default_eplb.py
      Default implementation of open-source EPLB paper algorithm
    • policy_swift_balancer.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_random.py
      Random policy for basic testing
    • 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 上,以确保它们之间的负载均衡。分层负载均衡策略可在预填充阶段使用,此时专家并行规模较小。

全局负载均衡

在其他情况下,我们使用全局负载均衡策略,该策略不考虑专家组,在全局范围内复制专家,并将复制的专家打包到各个 NPU 上。该策略可在解码阶段使用,此时专家并行规模较大。

添加新的 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. 例如:

    class RandomLoadBalance(EplbPolicy):
        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
    
  2. To add a new EPLB algorithm, include the policy type and its corresponding implementation class in the policy_factory.py of PolicyFactory.

Add a New MoE Model

Implementation Guide for Model Integration

  1. Adapter File Modification
  2. Inherit or modify vllm_ascend/eplb/adaptor/vllm_adaptor.py
  3. Add processing logic for key parameters:
    • num_dense_layers
    • global_expert_num
    • num_roe_layers
  4. Ensure parameter synchronization in the model_register function.

    For example:

    Modify vllm_adaptor.py of __init__ to add a new moe model eplb params:

       if self.model.config.model_type == "qwen3_moe":
        self.num_dense_layers = 0
        self.global_expert_num = self.model.config.num_experts
    

    Modify vllm_adaptor.py of model_register to register eplb params for new moe model:

        if config.model_type == "qwen3_moe":
            model.num_moe_layers = config.num_hidden_layers
    
  5. MoE Feature Integration

  6. Extend vllm_ascend/eplb/utils.py with MoE-specific methods
  7. Implement required functionality for expert routing or weight management

  8. Registration Logic Update

  9. Add patch logic within the model_register function
  10. Maintain backward compatibility with existing model types

  11. Validation & Testing

  12. Verify parameter consistency across layers
  13. Test cross-device communication for expert tables
  14. Benchmark against baseline implementations (e.g., Qwen3-MoE)

Key Implementation Notes:

  • Preserve existing interface contracts in abstract classes
  • Use decorators for non-intrusive patch integration
  • Leverage eplb_utils.py for shared expert mapping operations

DFX

Parameter Validation

Integer Parameters

All integer input parameters must explicitly specify their maximum and minimum values and be subject to valid value validation. For example, expert_heat_collection_interval 必须大于 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 be less than or equal to 0.")
        if iterations > sys.maxsize:
            raise ValueError(
                f"The {iterations} can not be larger 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 does 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 块处理函数体,指定捕获的异常类型和失败处理方式(例如记录异常或返回失败状态)。

一致性

专家映射

The expert map must be globally unique during initialization and update. In a multi-node scenario during initialization, distributed communication should be used to verify the consistency of expert maps across each rank. If they are inconsistent, the user should be notified of which ranks have inconsistent maps. During the update process, if only a few layers or the expert table of a certain rank has been changed, the updated expert table must be synchronized with the EPLB's context to ensure global consistency.

专家权重

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

限制

Before using EPLB, start the script and add export DYNAMIC_EPLB="true". Before performing load data collection (or performance data collection), start the script and add export EXPERT_MAP_RECORD="true".