From c22633b03a9173ba43878928b1fabfc1063fdc4d Mon Sep 17 00:00:00 2001 From: yurekami Date: Sat, 3 Jan 2026 03:11:29 +0900 Subject: [PATCH] fix: prevent duplicate logical experts on same GPU MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When packing physical experts to GPUs, the `balanced_packing` function previously only considered load balancing, which could result in multiple physical replicas of the same logical expert being placed on the same GPU. This is inefficient and defeats the purpose of expert replication. Changes: - Add optional `logical_ids` parameter to `balanced_packing` function - When packing, prefer packs that don't already contain the same logical ID - Pass `phy2mlog` (physical-to-logical mapping) to balanced_packing in Step 3 - Add comprehensive test suite with regression test for issue #22 The fix maintains backward compatibility - existing code without the logical_ids parameter continues to work as before. Fixes #22 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- eplb.py | 27 ++++++-- test_eplb.py | 186 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 206 insertions(+), 7 deletions(-) create mode 100644 test_eplb.py diff --git a/eplb.py b/eplb.py index 26c3987..c74e8c0 100644 --- a/eplb.py +++ b/eplb.py @@ -1,8 +1,9 @@ -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. @@ -10,8 +11,10 @@ def balanced_packing(weight: torch.Tensor, num_packs: int) -> Tuple[torch.Tensor 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) diff --git a/test_eplb.py b/test_eplb.py new file mode 100644 index 0000000..31e6c22 --- /dev/null +++ b/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!")