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

fix baichuan2 13b fp16 (#11071)

Browse source
This commit is contained in:
Yishuo Wang 2024-05-20 11:21:20 +08:00 committed by GitHub
parent 7170dd9192
commit 4e97047d70
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
python/llm/src/ipex_llm/transformers/models/baichuan2.py
View file

@ -213,6 +213,7 @@ def baichuan_attention_forward_13b(
attn_weights = attn_weights / math.sqrt(self.head_dim) attn_weights = attn_weights / math.sqrt(self.head_dim)
if attention_mask is not None: if attention_mask is not None:
attn_weights = attn_weights + attention_mask attn_weights = attn_weights + attention_mask
attn_weights = attn_weights.to(query_states.dtype)
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1) attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
if use_quantize_kv and q_len == 1: if use_quantize_kv and q_len == 1:
import linear_q4_0 import linear_q4_0