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"]