fix

Author: shihaobai

lightllm/models/qwen2_vl/infer_struct.py

@@ -11,11 +11,9 @@ def __init__(self):
         self.position_sin = None
 
     def init_some_extra_state(self, model, input_ids: torch.Tensor):
-        if "rope_type" in model.config["rope_scaling"]:
-            rope_scaling_type = model.config["rope_scaling"]["rope_type"]
-        elif "type" in model.config["rope_scaling"]:
-            rope_scaling_type = model.config["rope_scaling"]["type"]
-        if rope_scaling_type == "default":
+        rope_scaling = model.config["rope_scaling"].get("rope_type", {})
+        self.rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", None))
+        if self.rope_type != "mrope":
             super().init_some_extra_state(model, input_ids)
             return
         InferStateInfo.init_some_extra_state(self, model, input_ids)

lightllm/models/qwen2_vl/layer_infer/transformer_layer_infer.py

@@ -14,17 +14,13 @@ class Qwen2VLTransformerLayerInfer(LlamaTransformerLayerInfer):
     def __init__(self, layer_num, network_config, mode=[]):
         super().__init__(layer_num, network_config, mode)
         self.mrope_section = network_config["rope_scaling"]["mrope_section"]
-        if "rope_type" in network_config["rope_scaling"]:
-            self.rope_scaling_type = network_config["rope_scaling"]["rope_type"]
-        elif "type" in network_config["rope_scaling"]:
-            self.rope_scaling_type = network_config["rope_scaling"]["type"]
         axis_map = []
         for i, n in enumerate(self.mrope_section * 2):
             axis_map += [i % 3] * n
         self.axis_map = torch.tensor(axis_map, dtype=torch.int32, device="cuda")
 
     def _get_qkv(self, input, infer_state, layer_weight):
-        if self.rope_scaling_type == "default":
+        if infer_state.rope_type != "mrope":
             return super()._get_qkv(input, infer_state, layer_weight)
         q = layer_weight.q_proj.mm(input)
         cache_kv = layer_weight.kv_proj.mm(input).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)