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

Reduce Mistral softmax memory only in low memory mode (#11775)

Browse source 
* Reduce Mistral softmax memory only in low memory mode
This commit is contained in:
Qiyuan Gong 2024-08-13 14:50:54 +08:00 committed by GitHub
parent aa861df066
commit a88c132e54
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 5 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
python/llm/src/ipex_llm/transformers/models/mistral.py
View file

@ -131,11 +131,11 @@ def compute_attn_outputs_weights(query_states, key_states, value_states, bsz, q_
)
attn_weights = attn_weights + attention_mask
if kv_seq_len >= 2048 or bsz >= 64:
# for memory considerations, do not upcast attention to fp32
# for long sequences or large batches
attn_weights = nn.functional.softmax(attn_weights, dim=-1)
if os.getenv("IPEX_LLM_LOW_MEM", '0').lower() in ('true', '1', 't'):
if kv_seq_len >= 2048 or bsz >= 64:
# for memory considerations, do not upcast attention to fp32
# for long sequences or large batches
attn_weights = nn.functional.softmax(attn_weights, dim=-1)
else:
# upcast attention to fp32
attn_weights = nn.functional.softmax(attn_weights, dim=-1,
@ -838,7 +838,6 @@ def mistral_attention_forward_4_36_quantized(
f" but is {attention_mask.size()}"
)
attn_weights = attn_weights + attention_mask
if kv_seq_len >= 2048 or bsz >= 64:
# for memory considerations, do not upcast attention to fp32
# for long sequences or large batches