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

refactor qwen2 forward to enable XPU (#10409)

Browse source 
* refactor awen2 forward to enable XPU

* Update qwen2.py
This commit is contained in:
Heyang Sun 2024-03-14 11:03:05 +08:00 committed by GitHub
parent f36224aac4
commit 7d29765092
2 changed files with 6 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
python/llm/src/bigdl/llm/transformers/convert.py
View file

@ -1075,6 +1075,7 @@ def _optimize_post(model, lightweight_bmm=False):
modeling_module_name = model.__class__.__module__
module = importlib.import_module(modeling_module_name)
from bigdl.llm.transformers.models.qwen2 import qwen2_model_forward
from bigdl.llm.transformers.models.qwen2 import qwen2_attention_forward
convert_forward(model,
module.Qwen2Model,
qwen2_model_forward)
@ -1084,16 +1085,9 @@ def _optimize_post(model, lightweight_bmm=False):
convert_forward(model,
module.Qwen2MLP,
llama_mlp_forward)
if model.device.type == 'cpu':
from bigdl.llm.transformers.models.qwen2 import qwen2_sdpa_attention_forward
convert_forward(model,
module.Qwen2SdpaAttention,
qwen2_sdpa_attention_forward)
else:
from bigdl.llm.transformers.models.qwen2 import qwen2_attention_forward
convert_forward(model,
module.Qwen2Attention,
qwen2_attention_forward)
convert_forward(model,
module.Qwen2Attention,
qwen2_attention_forward)
elif model.config.model_type == "aquila":
modeling_module_name = model.__class__.__module__
module = importlib.import_module(modeling_module_name)

2
python/llm/src/bigdl/llm/transformers/models/qwen2.py
View file

@ -106,6 +106,8 @@ def qwen2_attention_forward(
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
if use_quantize_kv_cache(self.q_proj, hidden_states):
forward_function = qwen2_attention_forward_quantized
elif hidden_states.device.type == "cpu":
forward_function = qwen2_sdpa_attention_forward
else:
forward_function = qwen2_attention_forward_origin
return forward_function(