fix cache position for pytorch engine (#2388)

InternLM · Aug 27, 2024 · c57b635 · c57b635
1 parent 97b880b
commit c57b635
Show file tree

Hide file tree

Showing 7 changed files with 23 additions and 19 deletions.
diff --git a/lmdeploy/pytorch/models/baichuan.py b/lmdeploy/pytorch/models/baichuan.py
@@ -67,6 +67,7 @@ def forward(
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
+        **kwargs
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
                Optional[Tuple[torch.Tensor]]]:
         """Rewrite of Attention.forward."""
@@ -186,6 +187,7 @@ def forward(
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
+        **kwargs
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
                Optional[Tuple[torch.Tensor]]]:
         """Rewrite of BaichuanAttention.forward."""

diff --git a/lmdeploy/pytorch/models/chatglm2.py b/lmdeploy/pytorch/models/chatglm2.py
@@ -213,15 +213,14 @@ def _contiguous_batching_forward(
 
         return output, kv_cache
 
-    def forward(
-        self,
-        hidden_states,
-        attention_mask,
-        rotary_pos_emb,
-        kv_cache=None,
-        use_cache=True,
-        output_attentions=False,
-    ):
+    def forward(self,
+                hidden_states,
+                attention_mask,
+                rotary_pos_emb,
+                kv_cache=None,
+                use_cache=True,
+                output_attentions=False,
+                **kwargs):
         return self._contiguous_batching_forward(
             hidden_states,
             rotary_pos_emb,

diff --git a/lmdeploy/pytorch/models/deepseek.py b/lmdeploy/pytorch/models/deepseek.py
@@ -140,6 +140,7 @@ def forward(
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
+        **kwargs
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
                Optional[Tuple[torch.Tensor]]]:
         """forward."""

diff --git a/lmdeploy/pytorch/models/falcon.py b/lmdeploy/pytorch/models/falcon.py
@@ -215,16 +215,15 @@ def __rotary_emb_fn(query_states, key_states, value_states):
         else:
             return output_tensor, layer_past
 
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        alibi: Optional[torch.Tensor],
-        attention_mask: torch.Tensor,
-        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        head_mask: Optional[torch.Tensor] = None,
-        use_cache: bool = False,
-        output_attentions: bool = False,
-    ):
+    def forward(self,
+                hidden_states: torch.Tensor,
+                alibi: Optional[torch.Tensor],
+                attention_mask: torch.Tensor,
+                layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+                head_mask: Optional[torch.Tensor] = None,
+                use_cache: bool = False,
+                output_attentions: bool = False,
+                **kwargs):
         return self._contiguous_batching_forward(hidden_states, alibi,
                                                  layer_past)
 

diff --git a/lmdeploy/pytorch/models/internlm.py b/lmdeploy/pytorch/models/internlm.py
@@ -123,6 +123,7 @@ def forward(
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
+        **kwargs
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
                Optional[Tuple[torch.Tensor]]]:
         """forward."""

diff --git a/lmdeploy/pytorch/models/phi3.py b/lmdeploy/pytorch/models/phi3.py
@@ -152,6 +152,7 @@ def forward(
         past_key_value: Optional[Cache] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
+        **kwargs
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
                Optional[Tuple[torch.Tensor]]]:
         """rewrite of forward."""

diff --git a/lmdeploy/pytorch/models/starcoder2.py b/lmdeploy/pytorch/models/starcoder2.py
@@ -170,6 +170,7 @@ def forward(
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
         output_attentions: bool = False,
         use_cache: bool = False,
+        **kwargs
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor],
                Optional[Tuple[torch.Tensor]]]:
         """forward."""