add ep-size args

sgl-project · Dec 2, 2024 · 3aa1e25 · 3aa1e25
1 parent 35db97c
commit 3aa1e25
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
@@ -89,6 +89,8 @@ class ServerArgs:
     # Data parallelism
     dp_size: int = 1
     load_balance_method: str = "round_robin"
+    # Expert parallelism
+    ep_size: int = 1
 
     # Multi-node distributed serving
     dist_init_addr: Optional[str] = None
@@ -199,6 +201,12 @@ def __post_init__(self):
                 "Data parallel size is adjusted to be the same as tensor parallel size. "
                 "Overlap schedule is disabled."
             )
+        # Expert parallelism
+        if self.enable_ep_moe:
+            self.ep_size = self.tp_size
+            logger.info(
+                f"EP MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]."
+            )
 
     @staticmethod
     def add_cli_args(parser: argparse.ArgumentParser):
@@ -493,6 +501,14 @@ def add_cli_args(parser: argparse.ArgumentParser):
                 "shortest_queue",
             ],
         )
+        # Expert parallelism
+        parser.add_argument(
+            "--expert-parallel-size",
+            "--ep-size",
+            type=int,
+            default=ServerArgs.ep_size,
+            help="The expert parallelism size.",
+        )
 
         # Multi-node distributed serving
         parser.add_argument(
@@ -727,6 +743,7 @@ def add_cli_args(parser: argparse.ArgumentParser):
     def from_cli_args(cls, args: argparse.Namespace):
         args.tp_size = args.tensor_parallel_size
         args.dp_size = args.data_parallel_size
+        args.ep_size = args.expert_parallel_size
         attrs = [attr.name for attr in dataclasses.fields(cls)]
         return cls(**{attr: getattr(args, attr) for attr in attrs})
 

diff --git a/test/srt/test_moe_ep.py b/test/srt/test_moe_ep.py
@@ -20,7 +20,14 @@ def setUpClass(cls):
             cls.model,
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=["--tp", "2", "--trust-remote-code", "--enable-ep-moe"],
+            other_args=[
+                "--trust-remote-code",
+                "--tp",
+                "2",
+                "--ep-size",
+                "2",
+                "--enable-ep-moe",
+            ],
         )
 
     @classmethod
@@ -62,9 +69,11 @@ def setUpClass(cls):
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
             other_args=[
+                "--trust-remote-code",
                 "--tp",
                 "2",
-                "--trust-remote-code",
+                "--ep-size",
+                "2",
                 "--enable-ep-moe",
                 "--quantization",
                 "fp8",