From 5fb46157e678841cb8047194d404b3b563cbff1f Mon Sep 17 00:00:00 2001 From: yurekami Date: Fri, 26 Dec 2025 00:23:37 +0900 Subject: [PATCH 1/2] feat: Add Python packaging support and improve module interface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add proper Python packaging support to enable pip installation: - Add pyproject.toml with modern build system configuration - Add module docstring describing the package - Add __version__ = "1.0.0" for version tracking - Expand __all__ to export additional utility functions This enables users to install the package via: pip install git+https://github.com/deepseek-ai/EPLB.git Or for local development: pip install -e . 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- eplb.py | 17 ++++++++++++++++- pyproject.toml | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) create mode 100644 pyproject.toml diff --git a/eplb.py b/eplb.py index 26c3987..33b1a8c 100644 --- a/eplb.py +++ b/eplb.py @@ -1,7 +1,16 @@ +""" +Expert Parallelism Load Balancer (EPLB) + +A load balancing algorithm for expert parallelism in MoE (Mixture of Experts) models. +Implements redundant expert strategy with hierarchical and global load balancing policies. +""" + from typing import Tuple import torch +__version__ = "1.0.0" + def balanced_packing(weight: torch.Tensor, num_packs: int) -> 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 @@ -161,4 +170,10 @@ def rebalance_experts(weight: torch.Tensor, num_replicas: int, num_groups: int, torch.arange(num_replicas, dtype=torch.int64, device=log2phy.device).expand(num_layers, -1)) return phy2log, log2phy, logcnt -__all__ = ['rebalance_experts'] +__all__ = [ + "__version__", + "rebalance_experts", + "rebalance_experts_hierarchical", + "balanced_packing", + "replicate_experts", +] diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..c3980c6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,39 @@ +[build-system] +requires = ["setuptools >= 61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "eplb" +version = "1.0.0" +description = "Expert Parallelism Load Balancer for MoE models" +readme = "README.md" +license = {text = "MIT"} +authors = [ + {name = "DeepSeek-AI", email = "service@deepseek.com"}, +] +keywords = ["deep-learning", "moe", "expert-parallelism", "load-balancing", "pytorch"] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +requires-python = ">=3.8" +dependencies = [ + "torch >= 1.10.0", +] + +[project.urls] +Homepage = "https://github.com/deepseek-ai/EPLB" +Repository = "https://github.com/deepseek-ai/EPLB" +Issues = "https://github.com/deepseek-ai/EPLB/issues" + +[tool.setuptools] +py-modules = ["eplb"] From 60aca8bcd70153c5988628acbeee15dcb73377d0 Mon Sep 17 00:00:00 2001 From: yurekami Date: Tue, 30 Dec 2025 01:57:54 +0900 Subject: [PATCH 2/2] fix: prevent duplicate expert IDs on same GPU (#22) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Modified balanced_packing() to accept an optional item_ids parameter that prevents items with the same ID from being placed in the same pack. When packing physical experts to GPUs in rebalance_experts_hierarchical(), pass phy2mlog as item_ids to ensure no two physical experts that map to the same logical expert end up on the same GPU. This fixes the issue where hot experts could be duplicated on the same rank, e.g., "dup expert layer 0 rank 16, [200, 226, 236, 236, 12, ...]" 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- eplb.py | 35 +++++++++++++++++++++++++++++------ 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/eplb.py b/eplb.py index 33b1a8c..44b75b3 100644 --- a/eplb.py +++ b/eplb.py @@ -11,7 +11,8 @@ __version__ = "1.0.0" -def balanced_packing(weight: torch.Tensor, num_packs: int) -> Tuple[torch.Tensor, torch.Tensor]: +def balanced_packing(weight: torch.Tensor, num_packs: int, + item_ids: 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. @@ -19,8 +20,9 @@ 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: + item_ids: [X, n], optional item IDs - items with same ID won't be placed in same pack + + Returns: pack_index: [X, n], the pack index of each item rank_in_pack: [X, n], the rank of the item in the pack """ @@ -36,17 +38,37 @@ def balanced_packing(weight: torch.Tensor, num_packs: int) -> Tuple[torch.Tensor indices = weight.float().sort(-1, descending=True).indices.cpu() 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) + + if item_ids is not None: + item_ids = item_ids.cpu() + for i in range(num_layers): pack_weights = [0] * num_packs pack_items = [0] * num_packs + pack_ids = [set() for _ in range(num_packs)] if item_ids 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__) + group_id = item_ids[i, group].item() if item_ids is not None else None + + # Find eligible packs (not full and doesn't contain this item_id) + if item_ids is not None: + eligible = [p for p in range(num_packs) + if pack_items[p] < groups_per_pack and group_id not in pack_ids[p]] + if not eligible: + # Fallback if no eligible pack (constraints may be unsatisfiable) + eligible = [p for p in range(num_packs) if pack_items[p] < groups_per_pack] + else: + eligible = [p for p in range(num_packs) if pack_items[p] < groups_per_pack] + + pack = min(eligible, 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_ids is not None: + pack_ids[pack].add(group_id) + return pack_index, rank_in_pack @@ -123,8 +145,9 @@ def inverse(perm: torch.Tensor) -> torch.Tensor: # Step 3: pack physical_experts to GPUs # [num_layers * num_nodes, num_physical_experts // num_nodes] + # Pass phy2mlog as item_ids to prevent duplicate logical experts on the same GPU 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, item_ids=phy2mlog) phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack pphy2phy = inverse(phy2pphy)