diff --git a/python/llm/src/ipex_llm/transformers/npu_model.py b/python/llm/src/ipex_llm/transformers/npu_model.py index 00df48a7..a941ae29 100644 --- a/python/llm/src/ipex_llm/transformers/npu_model.py +++ b/python/llm/src/ipex_llm/transformers/npu_model.py @@ -116,12 +116,9 @@ class _BaseAutoModelClass: try: # for intel_npu_acceleration_library >= 1.1.0 from intel_npu_acceleration_library.quantization import quantize_model - from intel_npu_acceleration_library.compiler import ( - apply_horizontal_fusion, create_npu_kernels - ) + from intel_npu_acceleration_library.compiler import create_npu_kernels with torch.no_grad(): optimize_llm(model) - apply_horizontal_fusion(model) if not qtype.is_floating_point: model = quantize_model(model, qtype) create_npu_kernels(model) diff --git a/python/llm/src/ipex_llm/transformers/npu_models/common.py b/python/llm/src/ipex_llm/transformers/npu_models/common.py new file mode 100644 index 00000000..bb08b1ab --- /dev/null +++ b/python/llm/src/ipex_llm/transformers/npu_models/common.py @@ -0,0 +1,32 @@ +# +# Copyright 2016 The BigDL Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +from typing import List + + +def merge_linear(linears: List[torch.nn.Linear]) -> torch.nn.Linear: + new_weight = torch.cat(list(linear.weight.data for linear in linears), dim=0) + if linears[0].bias is not None: + new_linear = torch.nn.Linear(0, 0, bias=True) + new_bias = torch.cat(list(linear.bias.data for linear in linears), dim=0) + new_linear.bias = torch.nn.Parameter(new_bias, requires_grad=False) + else: + new_linear = torch.nn.Linear(0, 0, bias=False) + new_linear.weight = torch.nn.Parameter(new_weight, requires_grad=False) + new_linear.in_features = new_weight.size(1) + new_linear.out_features = new_weight.size(0) + return new_linear diff --git a/python/llm/src/ipex_llm/transformers/npu_models/convert.py b/python/llm/src/ipex_llm/transformers/npu_models/convert.py index 20482057..ba091815 100644 --- a/python/llm/src/ipex_llm/transformers/npu_models/convert.py +++ b/python/llm/src/ipex_llm/transformers/npu_models/convert.py @@ -29,6 +29,11 @@ def optimize_llm(model: torch.nn.Module): if model.config.model_type == "llama": from ipex_llm.transformers.npu_models.llama import merge_qkv model.apply(merge_qkv) + from ipex_llm.transformers.npu_models.llama import merge_mlp + model.apply(merge_mlp) from ipex_llm.transformers.npu_models.llama import llama_attention_forward from transformers.models.llama.modeling_llama import LlamaAttention convert_forward(model, LlamaAttention, llama_attention_forward) + from ipex_llm.transformers.npu_models.llama import llama_mlp_forward + from transformers.models.llama.modeling_llama import LlamaMLP + convert_forward(model, LlamaMLP, llama_mlp_forward) diff --git a/python/llm/src/ipex_llm/transformers/npu_models/llama.py b/python/llm/src/ipex_llm/transformers/npu_models/llama.py index 65d5f324..99b793dd 100644 --- a/python/llm/src/ipex_llm/transformers/npu_models/llama.py +++ b/python/llm/src/ipex_llm/transformers/npu_models/llama.py @@ -36,35 +36,33 @@ from typing import Optional, Tuple from transformers.cache_utils import Cache import torch -from transformers.models.llama.modeling_llama import LlamaAttention, repeat_kv, apply_rotary_pos_emb +from transformers.models.llama.modeling_llama import repeat_kv, apply_rotary_pos_emb +from transformers.models.llama.modeling_llama import LlamaAttention, LlamaMLP + +from ipex_llm.transformers.npu_models.common import merge_linear def merge_qkv(module: torch.nn.Module): if isinstance(module, LlamaAttention): - new_weight = torch.cat([ - module.q_proj.weight.data, - module.k_proj.weight.data, - module.v_proj.weight.data, - ], dim=0) - - if module.q_proj.bias is not None: - qkv_proj = torch.nn.Linear(0, 0, bias=True) - new_bias = torch.cat([ - module.q_proj.bias.data, - module.k_proj.bias.data, - module.v_proj.bias.data, - ], dim=0) - qkv_proj.bias = torch.nn.Parameter(new_bias, requires_grad=False) - else: - qkv_proj = torch.nn.Linear(0, 0, bias=False) - qkv_proj.weight = torch.nn.Parameter(new_weight, requires_grad=False) - qkv_proj.in_features = new_weight.size(1) - qkv_proj.out_features = new_weight.size(0) + qkv_proj = merge_linear([ + module.q_proj, + module.k_proj, + module.v_proj, + ]) module.qkv_proj = qkv_proj - del module.q_proj, module.k_proj, module.v_proj +def merge_mlp(module: torch.nn.Module): + if isinstance(module, LlamaMLP): + gate_up_proj = merge_linear([ + module.gate_proj, + module.up_proj, + ]) + module.gate_up_proj = gate_up_proj + del module.gate_proj, module.up_proj + + def llama_attention_forward( self, hidden_states: torch.Tensor, @@ -121,3 +119,10 @@ def llama_attention_forward( attn_weights = None return attn_output, attn_weights, past_key_value + + +def llama_mlp_forward(self, x): + gate_up_proj = self.gate_up_proj(x) + gate_proj, up_proj = gate_up_proj.chunk(2, dim=-1) + down_proj = self.down_proj(self.act_fn(gate_proj) * up_proj) + return down_proj