Merge pull request #451 from OpenAccess-AI-Collective/eval-is-causal

is_causal fix for evals?
2023-08-21 10:43:46 -07:00
parent 58cf7e7fed a213d9972a
commit 5e9c6afa10
1 changed files with 1 additions and 1 deletions
--- a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
+++ b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
@@ -158,7 +158,7 @@ def flashattn_forward(
    else:
        # turn off FA causal mask after first inference autoregressive iteration
        # only on first autoregressive step q,k,v have same seqlen
-        is_causal = past_key_value is not None
+        is_causal = key_states.shape == query_states.shape

    if cu_seqlens is not None and max_seqlen is not None:
        # special handling using sample packing