handle cache_position kwarg in updated llama modeling

axolotl-ai-cloud · Feb 21, 2024 · 9bd96e2 · 9bd96e2
1 parent 3eb834d
commit 9bd96e2
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
@@ -688,6 +688,9 @@ def llama_model_forward(
     output_attentions: Optional[bool] = None,
     output_hidden_states: Optional[bool] = None,
     return_dict: Optional[bool] = None,
+    cache_position: Optional[  # pylint: disable=unused-argument
+        torch.LongTensor
+    ] = None,
 ) -> Union[Tuple, BaseModelOutputWithPast]:
     output_attentions = (
         output_attentions