fix: prevent duplicate logical experts on same GPU

[yurekami]

Summary

This PR fixes a bug where the expert packing algorithm could place multiple physical replicas of the same logical expert on the same GPU, which defeats the purpose of expert replication for load balancing.

Problem

As reported in #22, when distributing 256 logical experts + 32 redundant experts across 32 GPUs, the algorithm could produce mappings like:

dup expert layer 0 rank 16, [200, 226, 236, 236, 12, 204, 156, 229, 184]

where expert 236 appears twice on the same GPU.

Root Cause

The balanced_packing function in Step 3 only considered load balancing when packing physical experts to GPUs. It did not check whether multiple physical experts mapping to the same logical expert would end up on the same GPU.

Solution

• Add optional logical_ids parameter to balanced_packing function
• When packing, prefer packs (GPUs) that don't already contain the same logical expert ID
• Pass the phy2mlog mapping to balanced_packing in Step 3 of rebalance_experts_hierarchical

The fix is backward compatible - existing code without the logical_ids parameter continues to work.

Changes

• eplb.py: Modified balanced_packing to accept and use logical ID constraints
• test_eplb.py: Added comprehensive test suite including:
  • README example validation
  • Regression test for issue eplb algorithm can generate phy2log with duplicate hot expert ids on same rank #22
  • High replication scenario test
  • Global load balancing test
  • Single replica per expert test

Test Results

README example test passed!
Issue #22 regression test passed!
High replication scenario test passed!
Global load balancing test passed!
Single replica per expert test passed!

All tests passed!

Test Plan

• All new tests pass
• README example produces valid output
• No duplicate logical experts on same GPU in stress tests
• Backward compatibility maintained (no API changes for existing callers)

Fixes #22

🤖 Generated with Claude Code