ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update llama.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-02-04 02:49:51 +11:00
commit 88a695da3a
1 changed files with 2 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

3
unsloth/models/llama.py
View file

@ -282,6 +282,7 @@ def LlamaAttention_fast_forward(
# Attention module
if (not HAS_FLASH_ATTENTION and attention_mask is None):
print(1)
# Xformers memory efficient attention
# Also has Flash Attention v2 dispatching
Q = Q.transpose(1, 2)
@ -309,7 +310,7 @@ def LlamaAttention_fast_forward(
V = V.transpose(1, 2)
A = flash_attn_func(Q, K, V, causal = True)
else:
print(attention_mask)
print(2)
# Grouped query attention
if n_groups != 1:
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, kv_seq_len, head_dim)