fix mem alloc (#627)

Co-authored-by: wangzaijun <[email protected]>
ModelTC · Nov 29, 2024 · 91d9e89 · 91d9e89
1 parent 01f8324
commit 91d9e89
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py
@@ -71,8 +71,15 @@ def __init__(self, kvargs):
         self._verify_must()
         self._verify_params()
         self._init_quant()
-        self._init_weights()
-        self._init_mem_manager()
+
+        # 更连续的显存分配可以有更好的性能
+        if self.max_total_token_num is None:
+            self._init_weights()
+            self._init_mem_manager()
+        else:
+            self._init_mem_manager()
+            self._init_weights()
+
         self._init_kv_move_buffer()
         self._check_mem_size()
         self._init_req_manager()

diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py
@@ -85,7 +85,7 @@ def init_model(self, kvargs):
             rank=self.tp_rank,
             world_size=self.world_size,
         )
-        
+
         from lightllm.distributed import set_custom_reduce
 
         set_custom_reduce()