From fcb884643bce8a83733ded21dc31031d1df3e120 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sun, 4 Feb 2024 02:02:51 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 5c16a050cc..8671d4e022 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -281,7 +281,7 @@ def LlamaAttention_fast_forward( past_key_value = (K, V) if use_cache else None # Attention module - if (not HAS_FLASH_ATTENTION): + if (not HAS_FLASH_ATTENTION and attention_mask is None): # Xformers memory efficient attention # Also has Flash Attention v2 dispatching Q = Q.transpose(1, 2) @@ -303,7 +303,7 @@ def LlamaAttention_fast_forward( A = xformers_attention(Q, K, V, attn_bias = causal_mask) A = A.view(bsz, q_len, n_heads, head_dim) - elif HAS_FLASH_ATTENTION: + elif HAS_FLASH_ATTENTION and attention_mask is None: Q = Q.transpose(1, 2) K = K.transpose(1, 2) V = V.transpose(1, 2)