Skip to content

[Cherry-Pick][Feature] Support redundant expert for eplb (#5918) #5923

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuanlehome merged 5 commits into from
Jan 9, 2026
Merged

[Cherry-Pick][Feature] Support redundant expert for eplb (#5918) #5923

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
20ef041
[Cherry-Pick] Support redundant expert for eplb
xiaoxiaohehe001 Jan 7, 2026
3615487
[Cherry-Pick] Support redundant expert for eplb
xiaoxiaohehe001 Jan 8, 2026
277de52
[Cherry-Pick] Support redundant expert for eplb
xiaoxiaohehe001 Jan 8, 2026
c07d52e
update
xiaoxiaohehe001 Jan 8, 2026
56f2636
update
xiaoxiaohehe001 Jan 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions custom_ops/gpu_ops/moe/ep_moe_expert_dispatch.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,11 @@
__VA_ARGS__ \
break; \
} \
case 7: { \
constexpr size_t NUM_EXPERTS_PER_RANK = 7; \
__VA_ARGS__ \
break; \
} \
case 8: { \
constexpr size_t NUM_EXPERTS_PER_RANK = 8; \
__VA_ARGS__ \
Expand All @@ -53,11 +58,26 @@
__VA_ARGS__ \
break; \
} \
case 10: { \
constexpr size_t NUM_EXPERTS_PER_RANK = 10; \
__VA_ARGS__ \
break; \
} \
case 16: { \
constexpr size_t NUM_EXPERTS_PER_RANK = 16; \
__VA_ARGS__ \
break; \
} \
case 17: { \
constexpr size_t NUM_EXPERTS_PER_RANK = 17; \
__VA_ARGS__ \
break; \
} \
case 20: { \
constexpr size_t NUM_EXPERTS_PER_RANK = 20; \
__VA_ARGS__ \
break; \
} \
case 32: { \
constexpr size_t NUM_EXPERTS_PER_RANK = 32; \
__VA_ARGS__ \
Expand Down
2 changes: 2 additions & 0 deletions fastdeploy/entrypoints/engine_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -593,6 +593,7 @@ async def rearrange_experts(self, request_dict: dict):
Returns:
tuple: response body, status code
"""
content, status_code = None, HTTPStatus.OK
eplb_config = self.config.eplb_config
if not eplb_config.enable_eplb:
content = {"code": 1, "msg": "redundant expert is disabled"}
Expand Down Expand Up @@ -695,6 +696,7 @@ async def get_per_expert_tokens_stats(self, request_dict: dict):
Returns:
tuple: response body, status code
"""
content, status_code = None, HTTPStatus.OK
eplb_config = self.config.eplb_config
if not eplb_config.enable_eplb:
content = {"code": 1, "msg": "redundant expert is disabled"}
Expand Down
2 changes: 1 addition & 1 deletion fastdeploy/model_executor/layers/backends/xpu/moe/ep.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor):
moe_topk=self.top_k,
apply_norm_weight=True, # apply_norm_weight
enable_softmax_top_k_fused=False,
redundant_ep_rank_num_plus_one=layer.fd_config.model_config.redundant_experts_num + 1,
redundant_ep_rank_num_plus_one=layer.fd_config.eplb_config.redundant_experts_num + 1,
)
else:
topk_idx, topk_weights = fastdeploy.model_executor.ops.xpu.moe_topk_select(
Expand Down
4 changes: 2 additions & 2 deletions fastdeploy/model_executor/layers/moe/ep.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -449,7 +449,7 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor):
expert_id_to_ep_rank_array=expert_id_to_ep_rank_array,
expert_in_rank_num_list=expert_in_rank_num_list,
tokens_per_expert_stats_list=tokens_per_expert_stats_list,
redundant_ep_rank_num_plus_one=layer.fd_config.model_config.redundant_experts_num + 1,
redundant_ep_rank_num_plus_one=layer.fd_config.eplb_config.redundant_experts_num + 1,
)
else:
topk_idx, topk_weights = fastdeploy.model_executor.ops.gpu.moe_redundant_topk_select(
Expand All @@ -461,7 +461,7 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor):
moe_topk=self.top_k,
apply_norm_weight=True,
enable_softmax_top_k_fused=False,
redundant_ep_rank_num_plus_one=layer.fd_config.model_config.redundant_experts_num + 1,
redundant_ep_rank_num_plus_one=layer.fd_config.eplb_config.redundant_experts_num + 1,
)
else:
if layer.topk_method == "noaux_tc":
Expand Down
2 changes: 1 addition & 1 deletion fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ def init_ep(self, layer: nn.Layer) -> None:
"num_max_dispatch_tokens_per_rank": layer.fd_config.model_config.num_max_dispatch_tokens_per_rank,
"ep_size": layer.ep_size,
"ep_rank": layer.ep_rank,
"redundant_experts_num": layer.fd_config.model_config.redundant_experts_num,
"redundant_experts_num": layer.fd_config.eplb_config.redundant_experts_num,
"ep_group": layer.fd_config.parallel_config.ep_group,
}

Expand Down
16 changes: 9 additions & 7 deletions fastdeploy/model_executor/layers/moe/moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -457,13 +457,18 @@ def load_experts_weight(
"""
logical_expert_ids = [
i
% (
self.fd_config.model_config.moe_num_experts[0]
if isinstance(self.fd_config.model_config.moe_num_experts, list)
else self.fd_config.model_config.moe_num_experts
)
for i in range(
self.expert_id_offset,
self.expert_id_offset + self.num_local_experts,
)
]
ep_rank_to_expert_id_list = [i for i in range(self.num_experts)]
if self.redundant_table_manger is not None and is_rearrange is True:
if self.redundant_table_manger is not None:
(
ep_rank_to_expert_id_list,
expert_id_to_ep_rank_array,
Expand All @@ -477,18 +482,15 @@ def load_experts_weight(
down_proj_weights = []
if isinstance(state_dict, list):
state_dict = dict(state_dict)
is_ffn_merged = (
up_gate_proj_expert_weight_key.format(logical_expert_ids[0] if is_rearrange else self.expert_id_offset)
in state_dict
)
is_ffn_merged = up_gate_proj_expert_weight_key.format(logical_expert_ids[0]) in state_dict
if is_ffn_merged:
for expert_idx in logical_expert_ids:
down_proj_expert_weight_key_name = down_proj_expert_weight_key.format(expert_idx)
up_gate_proj_expert_weight_key_name = up_gate_proj_expert_weight_key.format(expert_idx)
up_gate_proj_weights.append(
get_tensor(
(
state_dict.pop(up_gate_proj_expert_weight_key_name)
state_dict[up_gate_proj_expert_weight_key_name]
if up_gate_proj_expert_weight_key_name in state_dict
else up_gate_proj_expert_weight_key_name
),
Expand All @@ -498,7 +500,7 @@ def load_experts_weight(
down_proj_weights.append(
get_tensor(
(
state_dict.pop(down_proj_expert_weight_key_name)
state_dict[down_proj_expert_weight_key_name]
if down_proj_expert_weight_key_name in state_dict
else down_proj_expert_weight_key_name
),
Expand Down
4 changes: 4 additions & 0 deletions fastdeploy/model_executor/load_weight_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,8 +251,12 @@ def get_expert_ranges(fd_config):
"mtp_block" if getattr(fd_config.speculative_config, "model_type", "main") == "mtp" else "layers"
)

moe_num_experts = fd_config.model_config.moe_num_experts
if isinstance(moe_num_experts, list):
moe_num_experts = moe_num_experts[0]
for i in range(fd_config.model_config.moe_layer_start_index, fd_config.model_config.num_hidden_layers):
for j in get_expert_ranges(fd_config):
j = j % moe_num_experts
up_gate_proj_key = f"ernie.{prefix_layer_name}.{i}.mlp.experts.{j}.up_gate_proj.weight"
down_proj_key = f"ernie.{prefix_layer_name}.{i}.mlp.experts.{j}.down_proj.weight"

Expand Down
2 changes: 1 addition & 1 deletion fastdeploy/model_executor/models/ernie4_5_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -372,7 +372,7 @@ def __init__(
self.redundant_table_manger = RedundantExpertManger(
n_routed_experts=fd_config.model_config.moe_num_experts,
num_hidden_layers=fd_config.model_config.num_hidden_layers,
redundant_experts_num=fd_config.model_config.redundant_experts_num,
redundant_experts_num=fd_config.eplb_config.redundant_experts_num,
ep_size=fd_config.parallel_config.expert_parallel_size,
)

Expand Down
2 changes: 1 addition & 1 deletion fastdeploy/worker/experts_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ def __init__(
self.num_hidden_layers = num_hidden_layers

self.num_replicas = self.num_expert + self.redundant_experts_num
self.num_nodes = max(ep_size // 8, 1)
Copy link
Collaborator

@RichardWooSJTU RichardWooSJTU Jan 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个改动的目的是啥？那如果是单机呢

Copy link
Collaborator Author

@xiaoxiaohehe001 xiaoxiaohehe001 Jan 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这部分修改就是为了适配单机/多机，测试来看num_nodes=8能够保证初始化专家排布的正确性

self.num_nodes = max(ep_size // 8, 8)
self.num_gpus = ep_size
self.num_groups = 1

Expand Down
6 changes: 3 additions & 3 deletions fastdeploy/worker/worker_process.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -892,16 +892,17 @@ def initialize_fd_config(args, ranks: int = 1, local_rank: int = 0) -> FDConfig:
parallel_config = ParallelConfig(vars(args))
cache_config = CacheConfig(vars(args))
scheduler_config = SchedulerConfig(vars(args))
eplb_config = EPLBConfig(args.eplb_config)

parallel_config.tensor_parallel_rank = local_rank % parallel_config.tensor_parallel_size
parallel_config.data_parallel_rank = local_rank // parallel_config.tensor_parallel_size
# config for EP
if parallel_config.expert_parallel_size > 1:
expert_parallel_rank = int(local_rank % parallel_config.expert_parallel_size)
if isinstance(model_config.moe_num_experts, list):
num_experts = model_config.moe_num_experts[0]
num_experts = model_config.moe_num_experts[0] + eplb_config.redundant_experts_num
else:
num_experts = model_config.moe_num_experts
num_experts = model_config.moe_num_experts + eplb_config.redundant_experts_num
num_experts_per_rank = num_experts // parallel_config.expert_parallel_size
num_experts_start_offset = expert_parallel_rank * num_experts_per_rank
max_chips_per_node = 16 if current_platform.is_iluvatar() else 8
Expand All @@ -926,7 +927,6 @@ def initialize_fd_config(args, ranks: int = 1, local_rank: int = 0) -> FDConfig:
plas_attention_config = PlasAttentionConfig(args.plas_attention_config)

early_stop_config = EarlyStopConfig(args.early_stop_config)
eplb_config = EPLBConfig(args.eplb_config)

structured_outputs_config: StructuredOutputsConfig = StructuredOutputsConfig(args=vars(args))

Expand Down
Loading