is_causal fix for evals?

This commit is contained in:
Wing Lian
2023-08-21 10:36:26 -04:00
parent 58cf7e7fed
commit fbf49a4770

View File

@@ -155,6 +155,8 @@ def flashattn_forward(
# during training q,k,v always have same seqlen
assert key_states.shape == query_states.shape
is_causal = True
elif past_key_value is None:
is_causal = True
else:
# turn off FA causal mask after first inference autoregressive iteration
# only on first autoregressive step q,k,v have same seqlen