Skip to content

[Feature] Support redundant expert for eplb #5918

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
Jiang-Jia-Jun merged 6 commits into from
Jan 9, 2026
Merged

[Feature] Support redundant expert for eplb #5918

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
17b06bf
[BugFix] support redundant expert for eplb
xiaoxiaohehe001 Jan 7, 2026
b47e7e9
support redundant expert for eplb
xiaoxiaohehe001 Jan 7, 2026
40a336a
support redundant expert for eplb
xiaoxiaohehe001 Jan 8, 2026
947c38d
update
xiaoxiaohehe001 Jan 8, 2026
0dc3858
Merge branch 'develop' of https://github.com/PaddlePaddle/FastDeploy …
xiaoxiaohehe001 Jan 9, 2026
7617338
fix ci eplb
xiaoxiaohehe001 Jan 9, 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
10 changes: 10 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 @@ -48,6 +48,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 @@ -68,6 +73,11 @@
__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__ \
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 @@ -607,6 +607,7 @@ async def rearrange_experts(self, request_dict: dict):
Returns:
tuple: response body, status code
"""
content, status_code = None, HTTPStatus.OK
eplb_config = self.fd_config.eplb_config
if not eplb_config.enable_eplb:
content = {"code": 1, "msg": "redundant expert is disabled"}
Expand Down Expand Up @@ -709,6 +710,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.fd_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 @@ -277,7 +277,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 @@ -472,7 +472,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 @@ -484,7 +484,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 @@ -467,13 +467,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 @@ -487,18 +492,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
Copy link

Copilot AI Jan 9, 2026

Choose a reason for hiding this comment

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

移除了 is_rearrange 参数的检查，现在只要 redundant_table_manger 不为 None 就会执行重排逻辑。这个改动简化了逻辑，但需要确保所有调用此方法的地方都已更新，不再传递 is_rearrange 参数。此外，在第 495 行生成 is_ffn_merged 的检查中，也移除了对 is_rearrange 的依赖，确保逻辑的一致性。建议验证此改动不会影响非重排场景下的权重加载行为。

Suggested change
is_ffn_merged = up_gate_proj_expert_weight_key.format(logical_expert_ids[0]) in state_dict
if logical_expert_ids:
first_expert_key = up_gate_proj_expert_weight_key.format(logical_expert_ids[0])
is_ffn_merged = first_expert_key in state_dict
else:
# No local experts found, fall back to non-merged FFN loading path
is_ffn_merged = False

Copilot uses AI. Check for mistakes.
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 @@ -508,7 +510,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
5 changes: 5 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 @@ -252,8 +252,13 @@ 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):
# Map redundant expert IDs back to actual expert IDs for weight loading
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 @@ -400,7 +400,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)
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 @@ -927,6 +927,7 @@ 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
Expand All @@ -940,9 +941,9 @@ def initialize_fd_config(args, ranks: int = 1, local_rank: int = 0) -> FDConfig:
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
parallel_config.expert_parallel_rank = expert_parallel_rank
Expand All @@ -958,7 +959,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))
routing_replay_config = RoutingReplayConfig(args.routing_replay_config)
Expand Down
2 changes: 1 addition & 1 deletion tests/model_executor/test_ep.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -460,7 +460,7 @@ def get_ep_rank_to_expert_id_list_by_layer(self, _layer_idx):
top_k=2,
routed_scaling_factor=1.0,
gate_correction_bias=None,
fd_config=SimpleNamespace(model_config=SimpleNamespace(redundant_experts_num=0)),
fd_config=SimpleNamespace(eplb_config=SimpleNamespace(redundant_experts_num=0)),
)
gate_out = paddle.randn([1, 4], dtype="float32")

Expand Down
Loading