ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

LLM : Fix llama draft_model dtype error (#10005)

Browse source 
* fix llama draft_model dtype error

* updat
This commit is contained in:
Wang, Jian4 2024-01-26 10:59:48 +08:00 committed by GitHub
parent aae1870096
commit 98ea3459e5
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
python/llm/src/bigdl/llm/transformers/models/llama.py
View file

@ -729,7 +729,7 @@ def llama_attention_forward_4_36(
def native_sdp(query, key, value, attention_mask,
bsz, q_len, kv_seq_len, head_dim, num_heads):
attn_weights = torch.matmul(query,
attn_weights = torch.matmul(query.to(key.dtype),
key.transpose(2, 3)) / math.sqrt(head_dim)
attn_weights_size = (bsz, num_heads, q_len, kv_seq_len)