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

use sdp_causal to reduce internvl2-4b memory usage if set environment variable (#11953)

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

20
python/llm/src/ipex_llm/transformers/models/phi3.py
View file

@ -160,15 +160,17 @@ def attention_forward(
else: else:
attn_output = xe_addons.sdp(query_states, key_states, value_states, attn_output = xe_addons.sdp(query_states, key_states, value_states,
attention_mask) attention_mask)
# disable sdp_causal to avoid overflow for now elif (
# elif use_sdp_causal(q_len, kv_seq_len, self.head_dim, query_states, self.training): use_sdp_causal(q_len, kv_seq_len, self.head_dim, query_states, self.training)
# import xe_addons and os.environ.get("IPEX_LLM_LOW_MEM", "0") == "1"
# if isinstance(past_key_value, DynamicFp8Cache): ):
# attn_output = xe_addons.sdp_fp8_causal(query_states, key_states, import xe_addons
# value_states, attention_mask) if isinstance(past_key_value, DynamicFp8Cache):
# else: attn_output = xe_addons.sdp_fp8_causal(query_states, key_states,
# attn_output = xe_addons.sdp_causal(query_states, key_states, value_states, attention_mask)
# value_states, attention_mask) else:
attn_output = xe_addons.sdp_causal(query_states, key_states,
value_states, attention_mask)
else: else:
if use_quantizekv: if use_quantizekv:
key_states, value_states = restore_fp8_kv_cache(key_states, value_states, key_states, value_states = restore_fp8_kv_cache(key_states, value_states,