just pass args into torch checkpoint

axolotl-ai-cloud · Sep 26, 2023 · 8076bba · 8076bba
1 parent bb87204
commit 8076bba
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
@@ -520,9 +520,6 @@ def custom_forward(*inputs):
                     # None for past_key_value
                     return module(
                         *inputs,
-                        past_key_value,  # pylint: disable=(cell-var-from-loop)
-                        output_attentions,
-                        padding_mask=padding_mask,
                     )
 
                 return custom_forward
@@ -532,7 +529,10 @@ def custom_forward(*inputs):
                 hidden_states,
                 attention_mask,
                 position_ids,
+                past_key_value,
+                output_attentions,
                 None,
+                padding_mask,
                 cu_seqlens,
                 max_seqlen,
             )