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

optimize qwen2 gpu memory usage again (#11435)

Browse source
This commit is contained in:
Yishuo Wang 2024-06-26 16:52:29 +08:00 committed by GitHub
parent ab9f7f3ac5
commit 2a0f8087e3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 16 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
python/llm/src/ipex_llm/transformers/models/qwen2.py
View file

@ -357,6 +357,21 @@ def merge_qkv(module: torch.nn.Module):
del module.q_proj, module.k_proj, module.v_proj del module.q_proj, module.k_proj, module.v_proj
# Qwen2 uses pre-computed rope table to accelerate rope,
# original `cos_cached` and `sin_cached` are added by `register_buffer`,
# so they will move to xpu during `model.to('xpu')`.
# But gpu fuse kernel doesn't need this rope table, only cpu needs them,
# so delete them then add them with `=`, so that they will be pinned on CPU,
# this can save about 0.5GB gpu memory usage when running Qwen2
if hasattr(module.rotary_emb, "cos_cached"):
cos_cached = module.rotary_emb.cos_cached
del module.rotary_emb.cos_cached
module.rotary_emb.cos_cached = cos_cached
if hasattr(module.rotary_emb, "sin_cached"):
sin_cached = module.rotary_emb.sin_cached
del module.rotary_emb.sin_cached
module.rotary_emb.sin_cached = sin_cached
def padding_mlp(module: torch.nn.Module): def padding_mlp(module: torch.nn.Module):
# for qwen 1.5 14B # for qwen 1.5 14B
@ -433,6 +448,7 @@ def qwen2_attention_forward(
query_states, key_states) query_states, key_states)
else: else:
cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
cos, sin = cos.to(device), sin.to(device)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, query_states, key_states = apply_rotary_pos_emb(query_states, key_states,
cos, sin, position_ids) cos, sin, position_ids)