[LLM] support multi node deploy (#2708)

* [LLM] support multi node deploy

* Update engine.py

* fix bugs

* fix

* [LLM] support multi node deploy

* [LLM] support multi node deploy

---------

Co-authored-by: Jiang-Jia-Jun <163579578+Jiang-Jia-Jun@users.noreply.github.com>

Author: ltd0924

fastdeploy/cache_manager/cache_messager.py

@@ -37,6 +37,7 @@ class CacheMessager(object):
     def __init__(self,
                  splitwise_role,
                  transfer_protocol,
+                 pod_ip,
                  engine_worker_queue_port,
                  local_data_parallel_id,
                  gpu_cache_kvs,
@@ -69,7 +70,7 @@ def __init__(self,
         self.gpu_cache_kvs = gpu_cache_kvs
         self.rank = rank
         self.nranks = nranks
-        address = ('0.0.0.0', engine_worker_queue_port)
+        address = (pod_ip, engine_worker_queue_port)
         self.engine_worker_queue = EngineWorkerQueue(
             address=address,
             is_server=False,

fastdeploy/cache_manager/cache_transfer_manager.py

@@ -71,6 +71,10 @@ def parse_args():
                         type=int,
                         default=9923,
                         help="cache queue port")
+    parser.add_argument("--pod_ip",
+                        type=str,
+                        default="0.0.0.0",
+                        help="pod ip")
     parser.add_argument("--engine_worker_queue_port",
                         type=int,
                         default=9923,
@@ -144,7 +148,7 @@ def __init__(self, args):
         self.rank = rank
         self.device = device
 
-        address = ('0.0.0.0', args.cache_queue_port)
+        address = (args.pod_ip, args.cache_queue_port)
         self.cache_task_queue = EngineCacheQueue(
             address=address,
             is_server=False,
@@ -236,6 +240,7 @@ def __init__(self, args):
             self.cache_messager = CacheMessager(
                 splitwise_role=args.splitwise_role,
                 transfer_protocol=args.protocol,
+                pod_ip=args.pod_ip,
                 engine_worker_queue_port=args.engine_worker_queue_port,
                 local_data_parallel_id=args.local_data_parallel_id,
                 gpu_cache_kvs=self.gpu_cache_kvs,

fastdeploy/cache_manager/prefix_cache_manager.py

@@ -109,7 +109,7 @@ def __init__(self,
 
 
     def launch_cache_manager(self, cache_config, tensor_parallel_size, \
-                    device_ids, engine_worker_queue_port, pid_suffix):
+                    device_ids, pod_ip, engine_worker_queue_port, pid_suffix):
         """
         launch_cache_manager function used to initialize the cache manager.
         """
@@ -123,7 +123,7 @@ def launch_cache_manager(self, cache_config, tensor_parallel_size, \
             create=True)
 
         self.cache_task_queue = EngineCacheQueue(
-            address=('127.0.0.1', cache_config.cache_queue_port),
+            address=(pod_ip, cache_config.cache_queue_port),
             authkey=b'cache_queue_service',
             is_server=False,
             num_client=tensor_parallel_size,
@@ -166,6 +166,7 @@ def launch_cache_manager(self, cache_config, tensor_parallel_size, \
                 f" --cache_dtype {cache_config.cache_dtype}" +
                 f" --cache_queue_port {cache_config.cache_queue_port}" +
                 f" --enable_splitwise {int(self.enable_splitwise)}" +
+                f" --pod_ip {pod_ip}" +
                 f" --engine_worker_queue_port {engine_worker_queue_port}" +
                 f" --num_gpu_blocks {cache_config.total_block_num}" +
                 f" --num_cpu_blocks {cache_config.num_cpu_blocks}" +

fastdeploy/engine/args_utils.py

@@ -122,10 +122,7 @@ class EngineArgs:
     """
     Ratio of tokens to process in a block.
     """
-    nnode: int = 1
-    """
-    Number of nodes in the cluster.
-    """
+
     pod_ips: Optional[List[str]] = None
     """
     List of IP addresses for nodes in the cluster.
@@ -485,10 +482,7 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             default=EngineArgs.pod_ips,
             help=
             "List of IP addresses for nodes in the cluster (comma-separated).")
-        system_group.add_argument("--nnode",
-                                  type=int,
-                                  default=EngineArgs.nnode,
-                                  help="Number of nodes in the cluster.")
+
 
         # Performance tuning parameters group
         perf_group = parser.add_argument_group("Performance Tuning")
@@ -773,7 +767,6 @@ def create_engine_config(self) -> Config:
             max_num_seqs=self.max_num_seqs,
             speculative_config=speculative_cfg,
             max_num_batched_tokens=self.max_num_batched_tokens,
-            nnode=self.nnode,
             pod_ips=self.pod_ips,
             use_warmup=self.use_warmup,
             engine_worker_queue_port=self.engine_worker_queue_port,

fastdeploy/engine/config.py

@@ -505,7 +505,6 @@ def __init__(
         model_name_or_path: str = None,
         tokenizer: str = None,
         tensor_parallel_size: int = 8,
-        nnode: int = 1,
         max_model_len: int = 8192,
         max_num_seqs: int = 8,
         max_num_batched_tokens: Optional[int] = None,
@@ -539,7 +538,6 @@ def __init__(
             model_name_or_path (str): Model directory path or model name.
             tokenizer (str): Default is the model.
             tensor_parallel_size (int): Tensor parallel size. Default is 8.
-            nnode (int): Number of nodes. Default is 1.
             max_model_len (int): Maximum model length. Default is 8192.
             max_num_seqs (int): Maximum number of sequences. Default is 8.
             max_num_batched_tokens (Optional[int]): Maximum number of batched tokens. Default is None.
@@ -565,7 +563,6 @@ def __init__(
         self.tokenizer = tokenizer
         self.max_num_batched_tokens = max_num_batched_tokens
         self.tensor_parallel_size = tensor_parallel_size
-        self.nnode = nnode
         self.pod_ips = pod_ips
         self.max_model_len = max_model_len
         self.max_num_seqs = max_num_seqs
@@ -585,12 +582,15 @@ def __init__(
         self.max_capture_batch_size = max_capture_batch_size
         self.guided_decoding_backend = guided_decoding_backend
         self.disable_any_whitespace = disable_any_whitespace
+        self.is_master = True
+        self._str_to_list("innode_prefill_ports", int)
+        self._str_to_list("pod_ips", str)
 
-        if self.innode_prefill_ports is not None:
-            if not isinstance(self.innode_prefill_ports, list):
-                ports = str(self.innode_prefill_ports).split(',')
-                self.innode_prefill_ports = [int(port) for port in ports]
-
+        if self.pod_ips is None:
+            self.nnode = 1
+        else:
+            self.nnode = len(self.pod_ips)
+        
         assert self.splitwise_role in ["mixed", "prefill", "decode"]
 
         # TODO
@@ -609,14 +609,15 @@ def __init__(
 
         num_ranks = self.tensor_parallel_size * self.parallel_config.expert_parallel_size
         if num_ranks > 8:
-            local_num_ranks = 8
-            self.nnode = ceil_div(num_ranks, local_num_ranks)
+            self.worker_num_per_node = 8
+            nnode = ceil_div(num_ranks, self.worker_num_per_node)
+            assert nnode == self.nnode, \
+                f"nnode: {nnode}, but got {self.nnode}"
         else:
-            local_num_ranks = num_ranks
+            self.worker_num_per_node = num_ranks
 
         self.engine_worker_queue_port = engine_worker_queue_port
-        self.device_ids = ",".join([str(i) for i in range(min((self.tensor_parallel_size * \
-                                        self.parallel_config.expert_parallel_size), 8))])
+        self.device_ids = ",".join([str(i) for i in range(self.worker_num_per_node)])
         self.device_ids = os.getenv("CUDA_VISIBLE_DEVICES", self.device_ids)
 
         self.read_from_config()
@@ -628,16 +629,21 @@ def postprocess(self):
         """
         calculate some parameters
         """
-        total_rank = self.tensor_parallel_size * self.parallel_config.expert_parallel_size
-        assert self.device_ids.split(',').__len__() == min(total_rank, 8), \
-        f"invalid CUDA_VISIBLE_DEVICES, should be equal to {min(total_rank, 8)}"
+        assert self.device_ids.split(',').__len__() == self.worker_num_per_node, \
+        f"invalid CUDA_VISIBLE_DEVICES, should be equal to {self.worker_num_per_node}"
+
+        assert self.worker_num_per_node % self.tensor_parallel_size == 0, \
+            f"tensor_parallel_size: {self.tensor_parallel_size} should be divisible by worker_num_per_node: {self.worker_num_per_node}"
         self.local_device_ids = self.device_ids.split(
             ',')[:self.tensor_parallel_size]
-        assert self.tensor_parallel_size % self.nnode == 0, \
-        f"tensor_parallel_size: {self.tensor_parallel_size} should be divisible by nnode: {self.nnode}"
-        self.worker_num_per_node = total_rank // self.nnode
+
         self.host_ip = get_host_ip()
 
+        if self.pod_ips is None:
+            self.pod_ips = ["0.0.0.0"]
+        elif self.host_ip != self.pod_ips[0]:
+            self.is_master = False
+
         import paddle
         self.paddle_commit_id = paddle.version.commit
 
@@ -808,5 +814,16 @@ def reset_value(cls, value_name, key):
                     "return_full_hidden_states")
         reset_value(self.cache_config, "cache_dtype", "infer_model_dtype")
 
+    def _check_master(self):
+        return self.is_master
+    
+    def _str_to_list(self, attr_name, default_type):
+        if hasattr(self, attr_name):
+            val = getattr(self, attr_name)
+            if type(val) is str:
+                setattr(self, attr_name, [default_type(i) for i in val.split(",")])
+            else:
+                setattr(self, attr_name, val)
+
     def __str__(self) -> str:
         return json.dumps(self.__dict__, indent=4)

fastdeploy/engine/engine.py

@@ -98,30 +98,7 @@ def __init__(self, cfg):
                                                  cfg.mm_processor_kwargs,
                                                  cfg.enable_mm)
 
-        address = ('0.0.0.0', self.cfg.engine_worker_queue_port)
-        self.engine_worker_queue_server = EngineWorkerQueue(
-            address=address,
-            is_server=True,
-            num_client=self.cfg.tensor_parallel_size,
-            local_data_parallel_size=self.cfg.parallel_config.
-            data_parallel_size)
-
-        self.engine_worker_queue = EngineWorkerQueue(
-            address=address,
-            is_server=False,
-            num_client=self.cfg.tensor_parallel_size,
-            client_id=0,
-            local_data_parallel_id=0)
-
-        if self.cfg.cache_config.enable_prefix_caching or self.cfg.splitwise_role != 'mixed':
-            self.cache_task_queue = EngineCacheQueue(
-                address=('127.0.0.1', self.cfg.cache_config.cache_queue_port),
-                authkey=b'cache_queue_service',
-                is_server=True,
-                num_client=self.cfg.tensor_parallel_size,
-                client_id=-1,
-                local_data_parallel_size=self.cfg.parallel_config.
-                data_parallel_size)
+        self.start_queue_service()
 
         self.resource_manager = ResourceManager(cfg.max_num_seqs, cfg,
                                                 cfg.tensor_parallel_size,
@@ -198,9 +175,12 @@ def start(self, api_server_pid=None):
                 or self.cfg.splitwise_role != "mixed"):
             device_ids = self.cfg.device_ids.split(",")
             self.cache_manager_processes = self.resource_manager.cache_manager.launch_cache_manager(
-                self.cfg.cache_config, self.cfg.tensor_parallel_size,
-                device_ids, self.cfg.engine_worker_queue_port,
-                self.ipc_signal_suffix)
+                cache_config=self.cfg.cache_config,
+                tensor_parallel_size=self.cfg.tensor_parallel_size,
+                device_ids=device_ids,
+                pod_ip=self.cfg.pod_ips[0],
+                engine_worker_queue_port=self.cfg.engine_worker_queue_port,
+                pid_suffix=self.ipc_signal_suffix)
 
         self.worker_proc = self._start_worker_service()
         console_logger.info("Waitting worker processes ready...")
@@ -850,10 +830,7 @@ def _init_worker_signals(self):
         Initialize shared memory to indicate engine status
         """
         # worker_ready_signatensor_parallel_size
-        array_size = min(
-            8, self.cfg.tensor_parallel_size *
-            self.cfg.parallel_config.data_parallel_size)
-        worker_ready_signal_data = np.zeros(shape=[array_size], dtype=np.int32)
+        worker_ready_signal_data = np.zeros(shape=[self.cfg.worker_num_per_node], dtype=np.int32)
         self.worker_ready_signal = IPCSignal(name="worker_ready_signal",
                                              array=worker_ready_signal_data,
                                              dtype=np.int32,
@@ -889,7 +866,7 @@ def _init_worker_signals(self):
             create=True)
 
         # worker_live_signal 用于engine感知各worker进程是否存活，记录每个step 时间
-        worker_healthy_live_recorded_time_array = np.zeros(shape=[array_size],
+        worker_healthy_live_recorded_time_array = np.zeros(shape=[self.cfg.worker_num_per_node],
                                                            dtype=np.int32)
         self.worker_healthy_live_signal = IPCSignal(
             name="worker_healthy_live_signal",
@@ -899,7 +876,7 @@ def _init_worker_signals(self):
             create=True)
 
         if self.do_profile:
-            get_profile_block_num = np.zeros([array_size], dtype=np.int32)
+            get_profile_block_num = np.zeros([self.cfg.worker_num_per_node], dtype=np.int32)
             self.get_profile_block_num_signal = IPCSignal(
                 name="get_profile_block_num",
                 array=get_profile_block_num,
@@ -1028,13 +1005,15 @@ def _start_worker_service(self):
 
         arguments = (
             f" --nnodes {str(self.cfg.nnode)}"
+            f" --ips {','.join(self.cfg.pod_ips)}"
             f" --devices {self.cfg.device_ids} {py_script}"
             f" --max_num_seqs {self.cfg.max_num_seqs} --max_model_len {self.cfg.max_model_len}"
             f" --gpu_memory_utilization {self.cfg.cache_config.gpu_memory_utilization}"
             f" --model_name_or_path {str(self.cfg.model_name_or_path)}"
             f" --device_ids {self.cfg.device_ids}"
             f" --tensor_parallel_size {self.cfg.tensor_parallel_size}"
             f" --engine_worker_queue_port {str(self.cfg.engine_worker_queue_port)}"
+            f" --pod_ip {self.cfg.pod_ips[0]}"
             f" --total_block_num {self.cfg.cache_config.total_block_num}"
             f" --block_size {self.cfg.cache_config.block_size}"
             f" --enc_dec_block_num {self.cfg.cache_config.enc_dec_block_num}"
@@ -1171,10 +1150,12 @@ def _stop_profile(self):
         if self.cfg.cache_config.enable_prefix_caching or self.cfg.splitwise_role != "mixed":
             device_ids = self.cfg.device_ids.split(",")
             self.cache_manager_processes = self.resource_manager.cache_manager.launch_cache_manager(
-                self.cfg.cache_config, self.cfg.tensor_parallel_size,
-                device_ids, self.cfg.engine_worker_queue_port,
-                self.ipc_signal_suffix)
-
+                cache_config=self.cfg.cache_config,
+                tensor_parallel_size=self.cfg.tensor_parallel_size,
+                device_ids=device_ids,
+                pod_ip=self.cfg.pod_ips[0],
+                engine_worker_queue_port=self.cfg.engine_worker_queue_port,
+                pid_suffix=self.ipc_signal_suffix)
     def check_health(self, time_interval_threashold=30):
         """
         Check the health of the model server by checking whether all workers are alive.
@@ -1254,3 +1235,34 @@ def detect_thread():
         except Exception:
             pass
         return True
+
+    def start_queue_service(self):
+        """
+        start queue service for engine worker communication
+        """
+        address = (self.cfg.pod_ips[0], self.cfg.engine_worker_queue_port)
+        if self.cfg.host_ip == self.cfg.pod_ips[0] or self.cfg.pod_ips[0] == "0.0.0.0":
+            self.engine_worker_queue_server = EngineWorkerQueue(
+                address=address,
+                is_server=True,
+                num_client=self.cfg.tensor_parallel_size,
+                local_data_parallel_size=self.cfg.parallel_config.
+                data_parallel_size)
+
+            if self.cfg.cache_config.enable_prefix_caching or self.cfg.splitwise_role != 'mixed':
+                self.cache_task_queue = EngineCacheQueue(
+                    address=(self.cfg.pod_ips[0], self.cfg.cache_config.cache_queue_port),
+                    authkey=b'cache_queue_service',
+                    is_server=True,
+                    num_client=self.cfg.tensor_parallel_size,
+                    client_id=-1,
+                    local_data_parallel_size=self.cfg.parallel_config.
+                    data_parallel_size)
+
+
+        self.engine_worker_queue = EngineWorkerQueue(
+            address=address,
+            is_server=False,
+            num_client=self.cfg.tensor_parallel_size,
+            client_id=0,
+            local_data_parallel_id=0)

fastdeploy/engine/expert_service.py

@@ -65,7 +65,7 @@ def __init__(self, cfg, local_data_parallel_id):
 
         self.cfg.parallel_config.local_data_parallel_id = local_data_parallel_id
 
-        address = ('0.0.0.0', cfg.engine_worker_queue_port)
+        address = (cfg.pod_ips[0], cfg.engine_worker_queue_port)
         self.engine_worker_queue = EngineWorkerQueue(
             address=address,
             is_server=False,