is_causal fix for evals?
Browse files
src/axolotl/monkeypatch/llama_attn_hijack_flash.py
CHANGED
|
@@ -155,6 +155,8 @@ def flashattn_forward(
|
|
| 155 |
# during training q,k,v always have same seqlen
|
| 156 |
assert key_states.shape == query_states.shape
|
| 157 |
is_causal = True
|
|
|
|
|
|
|
| 158 |
else:
|
| 159 |
# turn off FA causal mask after first inference autoregressive iteration
|
| 160 |
# only on first autoregressive step q,k,v have same seqlen
|
|
|
|
| 155 |
# during training q,k,v always have same seqlen
|
| 156 |
assert key_states.shape == query_states.shape
|
| 157 |
is_causal = True
|
| 158 |
+
elif past_key_value is None:
|
| 159 |
+
is_causal = True
|
| 160 |
else:
|
| 161 |
# turn off FA causal mask after first inference autoregressive iteration
|
| 162 |
# only on first autoregressive step q,k,v have same seqlen
|