fix: fix a memleak

Author: niushengxiao

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -201,17 +201,24 @@ def init_model(self, kvargs):
             )
 
         # 用于协同读取 ShmObjsIOBuffer 中的请求信息的通信tensor和通信组对象。
-        self.node_broadcast_tensor = torch.tensor([0], dtype=torch.int32, device="cuda", requires_grad=False)
-        self.node_nccl_group = create_new_group_for_current_node("nccl")
+        # nccl频繁小数据通信会导致内存泄露, 这里需要换gloo
+        self.node_broadcast_tensor_cpu = torch.empty((1,), dtype=torch.int32, device="cpu", pin_memory=True)
+        self.node_gloo_group = create_new_group_for_current_node("gloo")
 
         # 用于在多节点tp模式下协同读取 ShmObjsIOBuffer 中的请求信息的通信tensor和通信组对象。
         if self.is_multinode_tp:
             self.multinode_tp_gather_item_tensor = torch.tensor([0], dtype=torch.int32, device="cuda")
             self.multinode_tp_all_gather_tensor = torch.tensor(
                 [0 for _ in range(self.global_world_size)], dtype=torch.int32, device="cuda", requires_grad=False
             )
-            self.multinode_tp_nccl_group = dist.new_group(
-                [rank for rank in range(self.global_world_size)], backend="nccl"
+            self.multinode_tp_gather_item_tensor_cpu = torch.empty(
+                (1,), dtype=torch.int32, device="cpu", pin_memory=True
+            )
+            self.multinode_tp_all_gather_tensor_cpu = torch.empty(
+                (self.global_world_size,), dtype=torch.int32, device="cpu", pin_memory=True
+            )
+            self.multinode_tp_gloo_group = dist.new_group(
+                [rank for rank in range(self.global_world_size)], backend="gloo"
             )
 
         if (
@@ -221,7 +228,7 @@ def init_model(self, kvargs):
             # 如果存在需要跨进程使用mem manger的特性，则将mem manager写入到 shm中，方便
             # 读取
             self.model.mem_manager.write_to_shm(req_manager=self.model.req_manager)
-            dist.barrier(group=self.node_nccl_group)
+            dist.barrier(group=self.node_gloo_group)
 
         self.init_custom()
 
@@ -362,28 +369,22 @@ def _try_read_new_reqs(self):
 
     def _try_read_new_reqs_normal(self):
         if self.is_master_in_node:
-            if self.shm_reqs_io_buffer.is_ready():
-                self.node_broadcast_tensor.fill_(1)
-            else:
-                self.node_broadcast_tensor.fill_(0)
+            self.node_broadcast_tensor_cpu[0] = 1 if self.shm_reqs_io_buffer.is_ready() else 0
 
         src_rank_id = self.args.node_rank * self.node_world_size
-        dist.broadcast(self.node_broadcast_tensor, src=src_rank_id, group=self.node_nccl_group, async_op=False)
-        new_buffer_is_ready = self.node_broadcast_tensor.detach().item()
+        dist.broadcast(self.node_broadcast_tensor_cpu, src=src_rank_id, group=self.node_gloo_group, async_op=False)
+        new_buffer_is_ready = int(self.node_broadcast_tensor_cpu[0].item())
         if new_buffer_is_ready:
             self._read_reqs_buffer_and_init_reqs()
 
         # nixl pd mode 从 shm_nixl_trans_io_buffer 读取分块传输的完成进度。
         if self.is_nixl_pd_mode:
             if self.is_master_in_node:
-                if self.shm_nixl_trans_io_buffer.is_ready():
-                    self.node_broadcast_tensor.fill_(1)
-                else:
-                    self.node_broadcast_tensor.fill_(0)
+                self.node_broadcast_tensor_cpu[0] = 1 if self.shm_nixl_trans_io_buffer.is_ready() else 0
 
             src_rank_id = self.args.node_rank * self.node_world_size
-            dist.broadcast(self.node_broadcast_tensor, src=src_rank_id, group=self.node_nccl_group, async_op=False)
-            new_buffer_is_ready = self.node_broadcast_tensor.detach().item()
+            dist.broadcast(self.node_broadcast_tensor_cpu, src=src_rank_id, group=self.node_gloo_group, async_op=False)
+            new_buffer_is_ready = int(self.node_broadcast_tensor_cpu[0].item())
             if new_buffer_is_ready:
                 self._read_nixl_trans_io_buffer_and_update_req_status()
         return
@@ -392,17 +393,14 @@ def _try_read_new_reqs_multinode_tp(self):
         """
         多节点tp模式下,需要协调所有rank的行为同步。
         """
-        if self.shm_reqs_io_buffer.is_ready():
-            self.multinode_tp_gather_item_tensor.fill_(1)
-        else:
-            self.multinode_tp_gather_item_tensor.fill_(0)
+        self.multinode_tp_gather_item_tensor_cpu[0] = 1 if self.shm_reqs_io_buffer.is_ready() else 0
         dist.all_gather_into_tensor(
-            self.multinode_tp_all_gather_tensor,
-            self.multinode_tp_gather_item_tensor,
-            group=self.multinode_tp_nccl_group,
+            self.multinode_tp_all_gather_tensor_cpu,
+            self.multinode_tp_gather_item_tensor_cpu,
+            group=self.multinode_tp_gloo_group,
             async_op=False,
         )
-        new_buffer_is_readys = self.multinode_tp_all_gather_tensor.detach().cpu().numpy()
+        new_buffer_is_readys = self.multinode_tp_all_gather_tensor_cpu.numpy()
         new_buffer_is_ready = np.all(new_buffer_is_readys == 1)
 
         if new_buffer_is_ready:

lightllm/server/router/model_infer/mode_backend/dp_backend/dp_shared_kv_trans.py

@@ -39,7 +39,7 @@ def fill_reqs_info(self, reqs: List[InferReq]):
         """
         填充请求的 kv 信息到共享内存中
         """
-        dist.barrier(group=self.backend.node_nccl_group)
+        dist.barrier(group=self.backend.node_gloo_group)
         if self.backend.is_master_in_dp:
             self.shared_req_infos.arr[0 : len(reqs), self.dp_rank_in_node, self._KV_LEN_INDEX] = [
                 req.cur_kv_len for req in reqs
@@ -57,7 +57,7 @@ def build_shared_kv_trans_tasks(
         """
         构建共享kv交换信息
         """
-        dist.barrier(group=self.backend.node_nccl_group)
+        dist.barrier(group=self.backend.node_gloo_group)
 
         trans_tasks: List[TransTask] = []
         rank_max_radix_cache_lens = np.max(