eplb algorithm can generate phy2log with duplicate hot expert ids on same rank

[SoonyangZhang]

I run eplb with collected topk traces. I found the algorithm can generate phy2log with duplicate hot expert ids on same rank.
I post the results here, there maybe some caution to handle for expert weight movement.
dup expert layer 0 rank 16, [200, 226, 236, 236, 12, 204, 156, 229, 184]

import torch
import eplb
def verify_same_expert_in_rank(phy2log, num_local_physical_experts):
    moe_layer, num_physical_experts = phy2log.shape
    world_size = num_physical_experts // num_local_physical_experts
    for layer_idx in range(moe_layer):
        indics = phy2log[layer_idx].tolist()
        for rank in range(world_size):
            expert_ids  = get_expert_in_rank(indics, num_local_physical_experts, rank)
            if len(expert_ids) != len(set(expert_ids)):
                print(f"dup expert layer {layer_idx} rank {rank}, {expert_ids}")

def test_eplb_single_layer():
    num_logical_experts = 256
    num_redundant_experts = 32
    num_replicas = num_logical_experts + num_redundant_experts
    num_groups = 16
    num_nodes = 2
    num_gpus = 32
    num_local_physical_experts = num_replicas // num_gpus
    device = "cpu"
    single_layer_load = [10545, 11846, 8542, 33117, 8385, 4892, 24052, 27875, 12287, 21604,
    31428, 19248, 13467, 12080, 5198, 1323, 19034, 28040, 44965, 3112,
    13323, 4808, 10167, 31079, 9812, 3159, 8383, 12858, 4414, 9934,
    23904, 25955, 13077, 8191, 8374, 4416, 24623, 20593, 1866, 11592,
    6743, 23520, 52598, 32672, 23038, 9723, 23399, 20137, 23853, 21483,
    11005, 8719, 12482, 17365, 20314, 21951, 12901, 7874, 22288, 11053,
    10730, 32446, 46416, 31024, 9464, 5708, 53472, 21800, 17393, 23875,
    7974, 4262, 29093, 3152, 76406, 35782, 11386, 10318, 7078, 39198,
    28862, 49373, 23481, 32155, 18591, 10249, 12558, 27202, 12870, 13441,
    25149, 9423, 21482, 5078, 10696, 6511, 26310, 11811, 12154, 13950,
    23652, 33239, 27797, 12445, 23949, 20423, 11732, 13073, 48033, 26085,
    34214, 11358, 26008, 12338, 31536, 5406, 6144, 12890, 8139, 27290,
    22171, 4439, 8982, 26922, 11461, 22704, 44095, 25014, 13137, 7628,
    6562, 18000, 9741, 10977, 8842, 11387, 34464, 10940, 6610, 9288,
    8782, 7979, 8467, 11106, 7738, 13468, 8217, 24875, 5315, 5996,
    8521, 35906, 39106, 21567, 22028, 10903, 8840, 25087, 7948, 22794,
    6346, 5121, 7601, 9692, 13697, 7753, 6755, 22514, 5902, 9048,
    21139, 26121, 10611, 1741, 16375, 8377, 5772, 9195, 14353, 7003,
    2359, 8007, 2383, 24437, 4113, 26870, 11119, 4029, 7255, 5515,
    6786, 8520, 11895, 7338, 7873, 26270, 25737, 58663, 11315, 9674,
    30632, 8648, 7468, 22430, 11192, 9860, 8022, 7877, 9884, 5390,
    13558, 9077, 5311, 6596, 30190, 12544, 17535, 22769, 11645, 34407,
    9020, 10113, 14553, 26278, 7948, 3933, 67904, 14050, 4083, 7298,
    34435, 29605, 9494, 9959, 26671, 40606, 36759, 28739, 22810, 1543,
    27385, 46296, 10938, 30118, 9053, 11687, 5958, 15211, 19456, 16559,
    7253, 9608, 8927, 22311, 7305, 11218]
    single_layer_expert_load = torch.tensor(single_layer_load, dtype=torch.int64, device=device)
    weight = single_layer_expert_load.unsqueeze(0)

    phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, 
        num_groups, num_nodes, num_gpus)
    verify_same_expert_in_rank(phy2log, num_local_physical_experts)

if __name__ == "__main__":
    test_eplb_single_layer()