DeepSeek/EPLB
1
0
Fork 0
mirror of https://github.com/deepseek-ai/EPLB synced 2025-04-28 18:01:42 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix inconsistent doc about the parameters.

Browse Source
This commit is contained in:
Haosdent Huang 2025-02-27 14:27:55 +08:00
parent f9bc62e841
commit fc36b96136
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
eplb.py
View File

@ -53,7 +53,7 @@ def replicate_experts(weight: torch.Tensor, num_phy: int) -> torch.Tensor:
Returns:
phy2log: [X, num_phy], logical expert id of each physical expert
rank: [X, num_phy], the duplica rank
rank: [X, num_phy], the replica rank
logcnt: [X, num_log], number of replicas for each logical expert
"""
n, num_log = weight.shape
@ -77,7 +77,10 @@ def rebalance_experts_hierarchical(weight: torch.Tensor, num_physical_experts: i
"""
Parameters:
weight: [num_moe_layers, num_logical_experts]
group_size: number of logical experts per group, used in group-limited routing
num_physical_experts: number of physical experts after replication
num_groups: number of expert groups
num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster
num_gpus: number of GPUs, must be a multiple of `num_nodes`
Returns:
physical_to_logical_map: [num_moe_layers, num_physical_experts]