fix: prevent duplicate logical experts on same GPU

eplb.py

@@ -1,17 +1,20 @@
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 
-def balanced_packing(weight: torch.Tensor, num_packs: int) -> Tuple[torch.Tensor, torch.Tensor]:
+def balanced_packing(weight: torch.Tensor, num_packs: int,
+                     logical_ids: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Pack n weighted objects to m packs, such that each bin contains exactly n/m objects and the weights of all packs
     are as balanced as possible.
 
     Parameters:
         weight: [X, n], the weight of each item
         num_packs: number of packs
-
-    Returns: 
+        logical_ids: [X, n], optional logical id per item; if provided, items with the same
+                     logical id will not be placed in the same pack when possible
+
+    Returns:
         pack_index: [X, n], the pack index of each item
         rank_in_pack: [X, n], the rank of the item in the pack
     """
@@ -25,19 +28,29 @@ def balanced_packing(weight: torch.Tensor, num_packs: int) -> Tuple[torch.Tensor
         return pack_index, rank_in_pack
 
     indices = weight.float().sort(-1, descending=True).indices.cpu()
+    logical_ids_cpu = logical_ids.cpu() if logical_ids is not None else None
     pack_index = torch.full_like(weight, fill_value=-1, dtype=torch.int64, device='cpu')
     rank_in_pack = torch.full_like(pack_index, fill_value=-1)
     for i in range(num_layers):
         pack_weights = [0] * num_packs
         pack_items = [0] * num_packs
+        pack_logicals = [set() for _ in range(num_packs)] if logical_ids_cpu is not None else None
         for group in indices[i]:
-            pack = min((i for i in range(num_packs) if pack_items[i] < groups_per_pack), 
-                       key=pack_weights.__getitem__)
+            logical_id = logical_ids_cpu[i, group].item() if logical_ids_cpu is not None else None
+            available_packs = [p for p in range(num_packs) if pack_items[p] < groups_per_pack]
+            if logical_id is not None:
+                # Prefer packs that don't already have this logical expert
+                constrained_packs = [p for p in available_packs if logical_id not in pack_logicals[p]]
+                if constrained_packs:
+                    available_packs = constrained_packs
+            pack = min(available_packs, key=pack_weights.__getitem__)
             assert pack_items[pack] < groups_per_pack
             pack_index[i, group] = pack
             rank_in_pack[i, group] = pack_items[pack]
             pack_weights[pack] += weight[i, group]
             pack_items[pack] += 1
+            if pack_logicals is not None:
+                pack_logicals[pack].add(logical_id)
     return pack_index, rank_in_pack
 
 
@@ -115,7 +128,7 @@ def inverse(perm: torch.Tensor) -> torch.Tensor:
     # Step 3: pack physical_experts to GPUs
     # [num_layers * num_nodes, num_physical_experts // num_nodes]
     tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog)
-    pack_index, rank_in_pack = balanced_packing(tokens_per_phy, num_gpus // num_nodes)
+    pack_index, rank_in_pack = balanced_packing(tokens_per_phy, num_gpus // num_nodes, logical_ids=phy2mlog)
     phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
     pphy2phy = inverse(phy2pphy)
 

test_eplb.py

@@ -0,0 +1,186 @@
+"""
+Tests for EPLB (Expert Parallelism Load Balancer)
+
+These tests verify the correctness of expert replication and placement,
+including regression tests for known issues.
+"""
+
+import torch
+import eplb
+
+
+def check_no_duplicate_experts_per_gpu(phy2log: torch.Tensor, num_gpus: int) -> bool:
+    """
+    Verify that no GPU has duplicate logical experts assigned to it.
+
+    Args:
+        phy2log: [num_layers, num_physical_experts], logical expert id for each physical expert
+        num_gpus: total number of GPUs
+
+    Returns:
+        True if no duplicates found, False otherwise
+    """
+    num_layers, num_phy = phy2log.shape
+    experts_per_gpu = num_phy // num_gpus
+
+    for layer in range(num_layers):
+        for gpu in range(num_gpus):
+            start = gpu * experts_per_gpu
+            end = start + experts_per_gpu
+            experts_on_gpu = phy2log[layer, start:end].tolist()
+            if len(experts_on_gpu) != len(set(experts_on_gpu)):
+                duplicates = [x for x in experts_on_gpu if experts_on_gpu.count(x) > 1]
+                print(f"Duplicate experts found! Layer {layer}, GPU {gpu}: {experts_on_gpu}")
+                print(f"Duplicates: {set(duplicates)}")
+                return False
+    return True
+
+
+def test_readme_example():
+    """Test the example from the README works correctly."""
+    weight = torch.tensor([[ 90, 132,  40,  61, 104, 165,  39,   4,  73,  56, 183,  86],
+                           [ 20, 107, 104,  64,  19, 197, 187, 157, 172,  86,  16,  27]])
+
+    num_replicas = 16
+    num_groups = 4
+    num_nodes = 2
+    num_gpus = 8
+
+    phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)
+
+    # Check shapes
+    assert phy2log.shape == (2, 16), f"Expected (2, 16), got {phy2log.shape}"
+    assert log2phy.shape[0] == 2 and log2phy.shape[1] == 12, f"Unexpected log2phy shape: {log2phy.shape}"
+    assert logcnt.shape == (2, 12), f"Expected (2, 12), got {logcnt.shape}"
+
+    # Check no duplicate experts per GPU
+    assert check_no_duplicate_experts_per_gpu(phy2log, num_gpus), "Found duplicate experts on same GPU!"
+
+    # Check all logical experts are covered
+    for layer in range(2):
+        logical_experts = set(phy2log[layer].tolist())
+        assert len(logical_experts) >= 12, f"Layer {layer}: Not all logical experts are assigned"
+
+    print("README example test passed!")
+
+
+def test_no_duplicate_experts_issue_22():
+    """
+    Regression test for issue #22: phy2log can generate duplicate hot expert ids on same rank.
+
+    This test uses a configuration similar to the bug report where 256 logical experts
+    with 32 redundant experts are distributed across 32 GPUs on 2 nodes.
+    """
+    # Simplified version of the bug report scenario
+    num_logical_experts = 256
+    num_redundant = 32
+    num_physical = num_logical_experts + num_redundant  # 288
+    num_gpus = 32
+    num_nodes = 2
+    num_groups = 8
+    num_layers = 4
+
+    # Create realistic load distribution with some experts being much hotter than others
+    torch.manual_seed(42)
+    weight = torch.rand(num_layers, num_logical_experts) * 100
+    # Make some experts much hotter (to trigger replication)
+    hot_experts = torch.randint(0, num_logical_experts, (20,))
+    weight[:, hot_experts] *= 10
+
+    phy2log, log2phy, logcnt = eplb.rebalance_experts(
+        weight, num_physical, num_groups, num_nodes, num_gpus
+    )
+
+    # Main assertion: no duplicate logical experts on the same GPU
+    assert check_no_duplicate_experts_per_gpu(phy2log, num_gpus), \
+        "Issue #22 regression: Found duplicate logical experts on same GPU!"
+
+    # Verify basic properties
+    assert phy2log.shape == (num_layers, num_physical)
+    assert logcnt.shape == (num_layers, num_logical_experts)
+
+    # Check that replication happened (some experts should have count > 1)
+    assert (logcnt > 1).any(), "Expected some experts to be replicated"
+
+    # Total replicas should match
+    assert logcnt.sum(dim=-1).tolist() == [num_physical] * num_layers
+
+    print("Issue #22 regression test passed!")
+
+
+def test_high_replication_scenario():
+    """
+    Test scenario with very high replication where duplicates are more likely.
+    """
+    num_logical_experts = 64
+    num_physical = 128  # 2x replication on average
+    num_gpus = 16
+    num_nodes = 2
+    num_groups = 4
+    num_layers = 2
+
+    # Create highly skewed load where some experts will be heavily replicated
+    weight = torch.ones(num_layers, num_logical_experts)
+    weight[:, :8] = 100  # First 8 experts are 100x hotter
+
+    phy2log, log2phy, logcnt = eplb.rebalance_experts(
+        weight, num_physical, num_groups, num_nodes, num_gpus
+    )
+
+    # Check no duplicates
+    assert check_no_duplicate_experts_per_gpu(phy2log, num_gpus), \
+        "Found duplicate logical experts in high replication scenario!"
+
+    # Check hot experts are replicated more
+    assert (logcnt[:, :8] > 1).all(), "Hot experts should be replicated"
+
+    print("High replication scenario test passed!")
+
+
+def test_global_load_balancing():
+    """Test the global load balancing policy (when num_groups % num_nodes != 0)."""
+    weight = torch.tensor([[10, 20, 30, 40, 50, 60],
+                           [60, 50, 40, 30, 20, 10]])
+
+    num_replicas = 12
+    num_groups = 3  # Not divisible by num_nodes=2
+    num_nodes = 2
+    num_gpus = 4
+
+    phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)
+
+    assert phy2log.shape == (2, 12)
+    assert check_no_duplicate_experts_per_gpu(phy2log, num_gpus), \
+        "Found duplicate experts in global load balancing mode!"
+
+    print("Global load balancing test passed!")
+
+
+def test_single_replica_per_expert():
+    """Test when there are no redundant experts (1:1 mapping)."""
+    weight = torch.rand(2, 16)
+
+    num_replicas = 16  # Same as logical experts
+    num_groups = 4
+    num_nodes = 2
+    num_gpus = 8
+
+    phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)
+
+    # All counts should be 1 (no replication)
+    assert (logcnt == 1).all(), "Expected no replication"
+
+    # Each logical expert appears exactly once
+    for layer in range(2):
+        assert sorted(phy2log[layer].tolist()) == list(range(16))
+
+    print("Single replica per expert test passed!")
+
+
+if __name__ == "__main__":
+    test_readme_example()
+    test_no_duplicate_experts_issue_22()
+    test_high_replication_scenario()
+    test_global_load_balancing()
+    test_single_replica_per_expert()
+    print("\nAll tests passed!")