动态专家负载均衡

动态专家负载均衡#

通用原理#

随着视觉生成模型向 DiT 架构演进，引入 MoE 机制以突破 Scaling Law 已成为行业共识。然而，DiT-MoE 庞大的参数规模迫使我们采用专家并行（EP）策略。与 LLM 不同，视觉数据的强空间局部性极易诱发特定专家过载，导致严重的计算负载不均。更进一步，扩散模型的去噪过程中专家激活分布呈现出显著的时序动态变化，这意味着传统的静态负载均衡策略在面对这种时空双重异构性时彻底失效。

技术特点#

本方案通过负载信息动态调整 Rank 上的专家权重以达到专家负载均衡，实现模型推理加速。方案具备以下特点：

无侵入式设计：全局同步点检查和权重更新位置可根据模型具体实现自行选择。
异步流水线处理：算法计算和专家权重拼接使用额外的线程和进程处理，最小化对主推理流程的影响。
三种 EP 模式：A2A（标准 all-to-all）、AG（all-gather）、EX（可控模式），通过 mode 参数选择。
与 CPU 卸载的互斥提醒：涉及 H2D 数据传输，与 CPU 卸载同时使用时可能存在带宽争抢，需自行调整执行时机。

接口和使用#

适配流程#

[!NOTE]说明为了最小程度的减少对主推理的影响，将算法和专家权重的拼接使用额外的线程和进程来处理。

启动 EPLB 算法进程。启动参数如下：

参数	默认值	说明
`world_size`	必填	EP 数
`expert_num`	必填	全局专家数量
`block_num`	必填	MoE 层数
`max_move`	—	EX 模式下最大移动专家数量
`redundant`	—	冗余专家数
`mode`	必填	A2A / AG / EX
`auth_key`	`secret_key`	默认读取环境变量 `EPLB_AUTH_KEY`

python -m mindiesd.eplb.eplb_scheduler \
    --world_size 2 \
    --host localhost \
    --port 50001 \
    --mode A2A

引入负载采集器和调度器，初始化后启动 worker 线程。

from mindiesd.eplb.dispatcher import DynamicDispatcher
from mindiesd.eplb.collector import ExpertLoadCollector
from mindiesd.eplb.task_manager import construct_expert_info_transfer_pool

model.init()

model.moe_module.block.expert_load_collector = ExpertLoadCollector(expert_num, lb_interval)
model.moe_module.block.dispatcher = DynamicDispatcher(expert_num, weight1, weight2, rank_in_group, ep_size)

if eplb_enabled:
    construct_expert_info_transfer_pool(
        module=model, rank_in_group=rank_in_group, device=device,
        ip=host, port=port, auth_key=auth_key
    )

model.forward()

AG 模式下需额外进行变换矩阵乘法。

if EP_AG and self.dispatcher.update_flag:
    expert_trans_tensor = self.dispatcher.get_expert_trans_tensor()
    trans_scores = torch.matmul(scores, expert_trans_tensor)

在 MoE 前向的 npu_moe_init_routing 之后、npu_grouped_matmul_finalize_routing 之前接入负载采集和权重替换。

expanded_tokens, expanded_row_idx, expanded_indices = torch_npu.npu_moe_init_routing(
    tokens, row_idx, indices, tokens.shape[0])

self.expert_load_collector.collect_expert_load(expanded_indices)
self.dispatcher.check_consistency()

if self.dispatcher.update_flag:
    weight1, weight2, local_expert_num, device_indices_map, \
        local_expert_indices_map, local_expert_list = \
        self.dispatcher.update_module_weight_and_map()
    self.weight1 = weight1
    self.weight2 = weight2
    self.local_expert_num = local_expert_num

tokens = torch_npu.npu_grouped_matmul_finalize_routing()

类说明#

ExpertLoadCollector#

from mindiesd.eplb.collector import ExpertLoadCollector

参数	类型	必选	默认值	说明
`expert_num`	`int`	是	-	全局专家数
`lb_interval`	`int`	否	`1`	EPLB 间隔步数

DynamicDispatcher#

from mindiesd.eplb.dispatcher import DynamicDispatcher

参数	类型	必选	默认值	说明
`expert_num`	`int`	是	-	全局专家数
`weight1`	`Tensor`	是	-	UP 权重
`weight2`	`Tensor`	是	-	DOWN 权重
`rank_in_group`	`int`	是	-	EP 通信组组内编号
`ep_size`	`int`	是	-	EP 数

construct_expert_info_transfer_pool#

from mindiesd.eplb.task_manager import construct_expert_info_transfer_pool

参数	类型	必选	默认值	说明
`module`	`Module`	是	-	初始化后的 model
`rank_in_group`	`int`	是	-	EP 通信组组内编号
`device`	`int`	是	-	rank 对应的 device 编号
`ip`	`str`	是	-	与服务端 ip 一致
`port`	`int`	是	-	与服务端 port 一致
`auth_key`	`str`	否	`secret_key`	默认读取环境变量 `EPLB_AUTH_KEY`