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

LLM: unify memory optimization env variables. (#11549)

Browse source 
* LLM: unify memory optimization env variables.

* fix comments.
This commit is contained in:
Cengguang Zhang 2024-07-11 11:01:28 +08:00 committed by GitHub
parent 51f2effb05
commit 70ab1a6f1a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 9 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
python/llm/src/ipex_llm/transformers/convert.py
View file

@ -327,9 +327,11 @@ def _replace_with_low_bit_linear(model, qtype, modules_to_not_convert=None,
optimize_lm_head = False optimize_lm_head = False
if is_lm_head(name, model_config, out_features): if is_lm_head(name, model_config, out_features):
model_type = getattr(model_config, "model_type", None) model_type = getattr(model_config, "model_type", None)
if model_type in ["gptj", "llama", "qwen2"] and \ if model_type in ["gptj", "llama", "qwen2"]:
os.environ.get("IPEX_LLM_LAST_LM_HEAD", None) == "1": if os.environ.get("IPEX_LLM_LAST_LM_HEAD", None) is not None:
optimize_lm_head = True optimize_lm_head = os.environ.get("IPEX_LLM_LAST_LM_HEAD", None) == "1"
elif os.environ.get("IPEX_LLM_LOW_MEM", None) is not None:
optimize_lm_head = os.environ.get("IPEX_LLM_LOW_MEM", None) == "1"
with init_empty_weights(): with init_empty_weights():
new_linear = None new_linear = None
is_gptq = is_gptq_linear(module) is_gptq = is_gptq_linear(module)

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

@ -286,6 +286,8 @@ def should_split_qkv_tensor(query_states, bsz, num_heads, q_len, kv_seq_len, out
if not output_attentions: if not output_attentions:
if os.environ.get("IPEX_LLM_SPLIT_QKV", None) is not None: if os.environ.get("IPEX_LLM_SPLIT_QKV", None) is not None:
return os.environ.get("IPEX_LLM_SPLIT_QKV", None) == "1" return os.environ.get("IPEX_LLM_SPLIT_QKV", None) == "1"
elif os.environ.get("IPEX_LLM_LOW_MEM", None) is not None:
return os.environ.get("IPEX_LLM_LOW_MEM", None) == "1"
elif query_states.dtype == torch.float16 and \ elif query_states.dtype == torch.float16 and \
query_states.shape[2] >= 6800: query_states.shape[2] >= 6800:
# split tensor for memory block limitation # split tensor for memory block limitation

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

@ -92,6 +92,8 @@ def should_split_qkv_tensor(query_states, bsz, num_heads, q_len, kv_seq_len, out
if not output_attentions: if not output_attentions:
if os.environ.get("IPEX_LLM_SPLIT_QKV", None) is not None: if os.environ.get("IPEX_LLM_SPLIT_QKV", None) is not None:
return os.environ.get("IPEX_LLM_SPLIT_QKV", None) == "1" return os.environ.get("IPEX_LLM_SPLIT_QKV", None) == "1"
elif os.environ.get("IPEX_LLM_LOW_MEM", None) is not None:
return os.environ.get("IPEX_LLM_LOW_MEM", None) == "1"
elif query_states.dtype == torch.float16 and \ elif query_states.dtype == torch.float16 and \
query_states.shape[2] >= 6300: query_states.shape[2] >= 6300:
# split tensor for memory block limitation # split tensor for memory block limitation