【Feature】 Refactor the code and add support for multi-output features like beam search. (#409)

Co-authored-by: wangzaijun <wangzaijun@sensetime.com>
Co-authored-by: shihaobai <42648726+shihaobai@users.noreply.github.com>
Co-authored-by: baishihao <baishihao@sensetime.com>

Author: 4 people

lightllm/common/basemodel/basemodel.py

@@ -38,15 +38,19 @@ def __init__(self, kvargs):
         self.weight_dir_ = kvargs["weight_dir"]
         self.max_total_token_num = kvargs["max_total_token_num"]
         self.load_way = kvargs.get("load_way", "HF")
-        self.mode = [m.replace('int4weight', 'w4a16').replace('int8weight', 'w8a16') for m in kvargs.get("mode", [])]
+        self.mode = [m.replace("int4weight", "w4a16").replace("int8weight", "w8a16") for m in kvargs.get("mode", [])]
         self.weight_dict = kvargs.get("weight_dict", None)
         self.finetune_config = kvargs.get("finetune_config", None)
         self.max_req_num = kvargs.get("max_req_num", 1000)
         self.max_seq_length = kvargs.get("max_seq_length", 1024 * 5)
-        self.return_all_prompt_logprobs = kvargs.get("return_all_prompt_logprobs", False)
+        # is_token_healing 和 return_all_prompt_logics 是有排斥关系的两个模式，只能单独有一个生效
+        # 主要是在prefill阶段返回多少个token的用于后续处理相关。
+        self.is_token_healing = kvargs.get("is_token_healing", False)
+        self.return_all_prompt_logics = kvargs.get("return_all_prompt_logics", False)
+        assert not (self.is_token_healing and self.return_all_prompt_logics), "can not be true in same time"
         self.use_dynamic_prompt_cache = kvargs.get("use_dynamic_prompt_cache", False)
         self.data_type = kvargs.get("data_type", "float16")
-        
+
         self._init_datatype()
         self._init_config()
         self._verify_must()
@@ -145,7 +149,7 @@ def _init_datatype(self):
         elif self.data_type in ["bf16", "bfloat16"]:
             self.data_type = torch.bfloat16
         elif self.data_type in ["fp32", "float32"]:
-            self.data_type =torch.float32
+            self.data_type = torch.float32
         else:
             raise ValueError(f"Unsupport datatype {self.data_type}!")
 
@@ -204,7 +208,8 @@ def _prefill(
     ):
         infer_state = self.infer_state_class()
         infer_state.is_prefill = True
-        infer_state.return_all_prompt_logprobs = self.return_all_prompt_logprobs
+        infer_state.is_token_healing = self.is_token_healing
+        infer_state.return_all_prompt_logics = self.return_all_prompt_logics
         infer_state.use_dynamic_prompt_cache = self.use_dynamic_prompt_cache
         infer_state.batch_size = batch_size
         infer_state.total_token_num = total_token_num

lightllm/common/basemodel/infer_struct.py

@@ -31,7 +31,8 @@ def __init__(self):
         self.kv_buffer = None
 
         self.is_splitfuse = False
-        self.return_all_prompt_logprobs = False
+        self.is_token_healing = False
+        self.return_all_prompt_logics = False
         self.use_dynamic_prompt_cache = False
         self.multimodal_params = None
 

lightllm/common/mem_manager.py

@@ -1,5 +1,7 @@
+import os
 import torch
 from lightllm.utils.log_utils import init_logger
+from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt
 
 logger = init_logger(__name__)
 
@@ -17,6 +19,13 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False
         self.mem_state = torch.zeros((size,), dtype=torch.int32, device="cuda")
         self.indexes = torch.arange(0, size, dtype=torch.long, device="cuda")
         self.can_use_mem_size = size
+        # 用共享内存进行共享，router 模块读取进行精确的调度估计, nccl port 作为一个单机中单实列的标记。防止冲突。
+        nccl_port = os.environ.get("_NCCL_PORT_", None)
+        assert nccl_port is not None
+        logger.info(f"mem manger get nccl port: {str(nccl_port)}")
+        self.shared_can_use_token_num = SharedInt(f"{str(nccl_port)}_mem_manger_can_use_token_num")
+
+        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
         self._init_buffers(size, dtype, head_num, head_dim, layer_num)
 
     def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
@@ -78,6 +87,7 @@ def add_refs(self, token_index: torch.Tensor):
         has_used_tokens = torch.count_nonzero(state).item()
         all_tokens = len(state)
         self.can_use_mem_size -= all_tokens - has_used_tokens
+        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
         self.mem_state[token_index] += 1
         return
 
@@ -89,11 +99,13 @@ def decrease_refs(self, token_index: torch.Tensor):
         used_tokens = torch.count_nonzero(state).item()
         all_tokens = len(state)
         self.can_use_mem_size += all_tokens - used_tokens
+        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
         return
 
     @torch.no_grad()
     def free_all(self):
         self.can_use_mem_size = len(self.mem_state)
+        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
         self.mem_state[:] = 0
 
     @torch.no_grad()
@@ -110,6 +122,7 @@ def resize_mem(self, new_size):
         self.mem_state = torch.zeros((size,), dtype=torch.int32, device="cuda")
         self.indexes = torch.arange(0, size, dtype=torch.long, device="cuda")
         self.can_use_mem_size = size
+        self.shared_can_use_token_num.set_value(self.can_use_mem_size)
         self._free_buffers()
         self._init_buffers(size, dtype, head_num, head_dim, layer_num)
         return

lightllm/common/req_manager.py

@@ -37,5 +37,3 @@ def free_token(self, free_token_index):
     def free_all(self):
         self.can_use_req_size = len(self.req_state)
         self.req_state[:] = 0
-    
-    

lightllm/models/llama/layer_infer/post_layer_infer.py

@@ -30,7 +30,9 @@ def _slice_get_last_input(self, input_embdings, infer_state: LlamaInferStateInfo
         if infer_state.is_splitfuse:
             # for SplitFuse
             batch_size = infer_state.batch_size
-            last_input = torch.empty((batch_size, self.embed_dim_), device=input_embdings.device, dtype=input_embdings.dtype)
+            last_input = torch.empty(
+                (batch_size, self.embed_dim_), device=input_embdings.device, dtype=input_embdings.dtype
+            )
             tmp_ = torch.cat(
                 [
                     torch.ones(infer_state.decode_req_num, dtype=torch.int32, device="cuda"),
@@ -42,16 +44,43 @@ def _slice_get_last_input(self, input_embdings, infer_state: LlamaInferStateInfo
             last_input[:, :] = input_embdings[last_index, :]
             return last_input, batch_size
 
-        if not infer_state.is_splitfuse and infer_state.is_prefill and not infer_state.return_all_prompt_logprobs:
+        if infer_state.is_prefill and infer_state.is_token_healing:
             batch_size = infer_state.batch_size
-            last_input = torch.empty((batch_size, self.embed_dim_), device=input_embdings.device, dtype=input_embdings.dtype)
+            b_seq_len_numpy = (infer_state.b_seq_len - infer_state.b_ready_cache_len).detach().cpu().numpy()
+            select_index = []
+            start_index = 0
+            select_token_num = 0
+            for cur_len in b_seq_len_numpy:
+                if cur_len == 1:
+                    select_index.append(start_index + cur_len - 1)
+                    start_index += cur_len
+                    select_token_num += 1
+                else:
+                    select_index.append(start_index + cur_len - 2)
+                    select_index.append(start_index + cur_len - 1)
+                    start_index += cur_len
+                    select_token_num += 2
+
+            last_index = torch.tensor(select_index, dtype=torch.long, device=input_embdings.device)
+            last_input = torch.empty(
+                (select_token_num, self.embed_dim_), device=input_embdings.device, dtype=input_embdings.dtype
+            )
+
+            last_input[:, :] = input_embdings[last_index, :]
+            return last_input, select_token_num
+
+        if not infer_state.is_splitfuse and infer_state.is_prefill and not infer_state.return_all_prompt_logics:
+            batch_size = infer_state.batch_size
+            last_input = torch.empty(
+                (batch_size, self.embed_dim_), device=input_embdings.device, dtype=input_embdings.dtype
+            )
             last_index = (
                 torch.cumsum(infer_state.b_seq_len - infer_state.b_ready_cache_len, dim=0, dtype=torch.long) - 1
             )
             last_input[:, :] = input_embdings[last_index, :]
             return last_input, batch_size
 
-        if not infer_state.is_splitfuse and infer_state.is_prefill and infer_state.return_all_prompt_logprobs:
+        if not infer_state.is_splitfuse and infer_state.is_prefill and infer_state.return_all_prompt_logics:
             total_tokens = infer_state.total_token_num
             return input_embdings, total_tokens
 
@@ -82,7 +111,9 @@ def token_forward(
         if self.world_size_ == 1:
             gather_data = logic_batch
         else:
-            gather_data = torch.empty((self.vocab_size_, token_num), device=logic_batch.device, dtype=input_embdings_dtype)
+            gather_data = torch.empty(
+                (self.vocab_size_, token_num), device=logic_batch.device, dtype=input_embdings_dtype
+            )
             split_indexes = np.linspace(0, self.vocab_size_, self.world_size_ + 1, dtype=np.int64)
             dist.all_gather(
                 [gather_data[split_indexes[i] : split_indexes[i + 1], :] for i in range(self.world_size_)],

lightllm/models/llava/llava_visual.py

@@ -7,49 +7,49 @@
 
 
 class LlavaVisionModel:
-
     def __init__(self):
         pass
 
     def load_model(self, weight_dir):
         config_file = os.path.join(weight_dir, "config.json")
         config = json.load(open(config_file))
-        self.select_layer = config.get('mm_vision_select_layer', -2)
-        self.select_feature = config.get('mm_vision_select_feature', 'patch')
+        self.select_layer = config.get("mm_vision_select_layer", -2)
+        self.select_feature = config.get("mm_vision_select_feature", "patch")
 
         # load clip vision model by cfg['mm_vision_tower']:
         #   huggingface_name or path_of_clip_relative_to_llava_model_dir
-        vision_path = config.get('mm_vision_tower', 'openai/clip-vit-large-patch14-336')
+        vision_path = config.get("mm_vision_tower", "openai/clip-vit-large-patch14-336")
         if isinstance(vision_path, list):
             vision_path = vision_path[0]
         if vision_path.startswith("./"):
             vision_path = os.path.join(weight_dir, vision_path)
 
         from transformers import CLIPVisionModel, CLIPImageProcessor
+
         self.image_processor = CLIPImageProcessor.from_pretrained(vision_path)
         self.vision_tower = CLIPVisionModel.from_pretrained(vision_path).half()
         self.vision_tower.requires_grad_(False)
-        self.device = torch.device('cpu')
+        self.device = torch.device("cpu")
 
         # load projector weights
         self.projector_weights = {}
         for f in os.listdir(weight_dir):
             if f.endswith(".bin"):
                 d = torch.load(os.path.join(weight_dir, f), "cpu")
                 for k, v in d.items():
-                    if 'model.mm_projector' in k:
+                    if "model.mm_projector" in k:
                         self.projector_weights[k] = v.half()
 
-        assert 'model.mm_projector.0.weight' in self.projector_weights
-        assert 'model.mm_projector.0.bias' in self.projector_weights
-        assert 'model.mm_projector.2.weight' in self.projector_weights
-        assert 'model.mm_projector.2.bias' in self.projector_weights
+        assert "model.mm_projector.0.weight" in self.projector_weights
+        assert "model.mm_projector.0.bias" in self.projector_weights
+        assert "model.mm_projector.2.weight" in self.projector_weights
+        assert "model.mm_projector.2.bias" in self.projector_weights
 
     def cuda(self):
         self.vision_tower = self.vision_tower.cuda()
         for k, v in self.projector_weights.items():
             self.projector_weights[k] = v.cuda()
-        self.device = torch.device('cuda')
+        self.device = torch.device("cuda")
         return self
 
     # batch images infer
@@ -58,7 +58,7 @@ def forward(self, x):
 
         x = self.vision_tower(x, output_hidden_states=True)
         x = x.hidden_states[self.select_layer]
-        if self.select_feature == 'patch':
+        if self.select_feature == "patch":
             x = x[:, 1:].contiguous()
         B, L, N = x.shape
         x = x.view(-1, N)
@@ -84,10 +84,12 @@ def encode(self, image_items: List[Union[str, Image.Image]]):
             if isinstance(item, Image.Image):
                 image = item
             elif item.startswith("http://") or item.startswith("https://"):
+                import requests
+
                 image = Image.open(requests.get(item, stream=True).raw)
             else:
                 image = Image.open(item)
             images.append(image.convert("RGB"))
 
-        images = self.image_processor.preprocess(images, return_tensors='pt')['pixel_values']
+        images = self.image_processor.preprocess(images, return_tensors="pt")["pixel_values"]
         return self.forward(images)

lightllm/models/starcoder_wquant/layer_infer/transformer_layer_infer.py

@@ -69,7 +69,7 @@ def _get_o(
         self, input, infer_state: StarcoderInferStateInfo, layer_weight: StarcoderTransformerLayerWeightQuantized
     ) -> torch.Tensor:
         o_output = self._wquant_matmul_for_o(
-            input.view(-1, self.embed_dim_), layer_weight.o_weight_, infer_state=infer_state, bias=layer_weight.o_bias_
+            input, layer_weight.o_weight_, infer_state=infer_state, bias=layer_weight.o_bias_
         )
         return o_output
 

lightllm/server/__init__.py

@@ -0,0 +1 @@
+from .router.token_load import TokenLoad