[Fix] Fix port error in PD disaggregation setting (#1175)

Author: yitianlian

slime/backends/sglang_utils/sglang_engine.py

@@ -91,7 +91,7 @@ def __init__(self, args, rank: int, worker_type: str = "regular"):
         self.rank = rank
         self.worker_type = worker_type
 
-    def init(self, dist_init_addr, port, nccl_port, host=None):
+    def init(self, dist_init_addr, port, nccl_port, host=None, disaggregation_bootstrap_port=None):
         self.router_ip = self.args.sglang_router_ip
         self.router_port = self.args.sglang_router_port
 
@@ -108,7 +108,14 @@ def init(self, dist_init_addr, port, nccl_port, host=None):
             dist_init_addr = f"[{ipv6_addr}]:{port_str}"
 
         server_args_dict, external_engine_need_check_fields = _compute_server_args(
-            self.args, self.rank, dist_init_addr, nccl_port, host, port, self.worker_type
+            self.args,
+            self.rank,
+            dist_init_addr,
+            nccl_port,
+            host,
+            port,
+            self.worker_type,
+            disaggregation_bootstrap_port,
         )
 
         self.node_rank = server_args_dict["node_rank"]
@@ -157,12 +164,15 @@ def _init_normal(self, server_args_dict):
                     f"http://{self.router_ip}:{self.router_port}/add_worker?url=http://{self.server_host}:{self.server_port}"
                 )
             else:
+                payload = {
+                    "url": f"http://{self.server_host}:{self.server_port}",
+                    "worker_type": self.worker_type,
+                }
+                if self.worker_type == "prefill":
+                    payload["bootstrap_port"] = server_args_dict["disaggregation_bootstrap_port"]
                 response = requests.post(
                     f"http://{self.router_ip}:{self.router_port}/workers",
-                    json={
-                        "url": f"http://{self.server_host}:{self.server_port}",
-                        "worker_type": self.worker_type,
-                    },
+                    json=payload,
                 )
             response.raise_for_status()
 
@@ -381,7 +391,16 @@ def stop_profile(self):
         return response
 
 
-def _compute_server_args(args, rank, dist_init_addr, nccl_port, host, port, worker_type: str = "regular"):
+def _compute_server_args(
+    args,
+    rank,
+    dist_init_addr,
+    nccl_port,
+    host,
+    port,
+    worker_type: str = "regular",
+    disaggregation_bootstrap_port: int | None = None,
+):
     nnodes = max(1, args.rollout_num_gpus_per_engine // args.num_gpus_per_node)
     node_rank = rank % nnodes
     kwargs = {
@@ -411,6 +430,10 @@ def _compute_server_args(args, rank, dist_init_addr, nccl_port, host, port, work
     if worker_type == "prefill":
         kwargs["disaggregation_mode"] = "prefill"
         kwargs["load_balance_method"] = "round_robin"
+        assert (
+            disaggregation_bootstrap_port is not None
+        ), "disaggregation_bootstrap_port must be set for prefill worker"
+        kwargs["disaggregation_bootstrap_port"] = disaggregation_bootstrap_port
     elif worker_type == "decode":
         kwargs["disaggregation_mode"] = "decode"
         kwargs["prefill_round_robin_balance"] = True
@@ -419,7 +442,6 @@ def _compute_server_args(args, rank, dist_init_addr, nccl_port, host, port, work
         kwargs["enable_return_routed_experts"] = True
     if args.fp16:
         kwargs["dtype"] = "float16"
-
     external_engine_need_check_fields = [k for k in kwargs.keys() if k not in _EXTERNAL_ENGINE_SKIP_CHECK_FIELDS]
 
     unused_keys = set(kwargs.keys())

slime/ray/rollout.py

@@ -435,6 +435,12 @@ def _allocate_rollout_engine_addr_and_ports_normal(*, args, num_engines, rollout
     )
     addr_and_ports = [{} for _ in range(num_engines)]
 
+    # Calculate prefill limit to identify prefill engines
+    prefill_limit = 0
+    if args.prefill_num_servers is not None:
+        num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node)
+        prefill_limit = args.prefill_num_servers * args.rollout_num_gpus_per_engine // num_gpu_per_engine
+
     visited_nodes = set()
     for rank, engine in rollout_engines:
         if rank // num_engines_per_node in visited_nodes:
@@ -469,9 +475,13 @@ def addr():
         get_addr, get_port = get_addr_and_ports(engine)
 
         for i in range(num_engines_on_this_node):
-            addr_and_ports[rank + i]["host"] = get_addr()
-            addr_and_ports[rank + i]["port"] = get_port()
-            addr_and_ports[rank + i]["nccl_port"] = get_port()
+            current_rank = rank + i
+            addr_and_ports[current_rank]["host"] = get_addr()
+            addr_and_ports[current_rank]["port"] = get_port()
+            addr_and_ports[current_rank]["nccl_port"] = get_port()
+
+            if args.prefill_num_servers is not None and current_rank < prefill_limit:
+                addr_and_ports[current_rank]["disaggregation_bootstrap_port"] = get_port()
 
         if args.rollout_num_gpus_per_engine > args.num_gpus_per_node:
             num_node_per_engine = args.rollout_num_gpus_per_engine // args.num_gpus_per_node

slime/utils/arguments.py

@@ -1581,6 +1581,9 @@ def slime_validate_args(args):
 
     if args.prefill_num_servers is not None:
         assert not args.use_fault_tolerance, "fault tolerance is not supported when prefill_num_servers is set."
+    assert not (
+        args.prefill_num_servers is not None and args.rollout_external
+    ), "prefill_num_servers cannot be set when rollout_external is set."
 
 
 def hf_validate_args(args, hf_config):