From 324bcb057e2453b614085f17be4cab74ef317a3f Mon Sep 17 00:00:00 2001 From: Yishuo Wang Date: Thu, 17 Oct 2024 14:45:09 +0800 Subject: [PATCH] refactor to reduce old rope usage (#12219) --- .../ipex_llm/transformers/models/aquila.py | 72 ++++------------ .../ipex_llm/transformers/models/decilm.py | 86 +++---------------- .../src/ipex_llm/transformers/models/llama.py | 4 +- .../ipex_llm/transformers/models/phixtral.py | 30 ------- 4 files changed, 31 insertions(+), 161 deletions(-) diff --git a/python/llm/src/ipex_llm/transformers/models/aquila.py b/python/llm/src/ipex_llm/transformers/models/aquila.py index 088c0fa5..b889c1c1 100644 --- a/python/llm/src/ipex_llm/transformers/models/aquila.py +++ b/python/llm/src/ipex_llm/transformers/models/aquila.py @@ -36,22 +36,14 @@ # limitations under the License. import math -from typing import List, Optional, Tuple, Union - import torch -import torch.utils.checkpoint -from torch import nn -from ipex_llm.transformers.models.utils import extend_kv_cache, init_kv_cache, \ - append_kv_cache, is_enough_kv_cache_room_4_31 -from ipex_llm.transformers.models.utils import apply_rotary_pos_emb -from ipex_llm.transformers.models.utils import apply_rotary_pos_emb_no_cache_xpu +from typing import Optional, Tuple +from ipex_llm.transformers.models.common import attention_softmax +from ipex_llm.transformers.models.utils import apply_rotary_pos_emb, should_use_fuse_rope +from ipex_llm.transformers.models.utils import update_past_key_value from ipex_llm.utils.common import log4Error -import os - -KV_CACHE_ALLOC_BLOCK_LENGTH = int(os.environ.get("KV_CACHE_ALLOC_BLOCK_LENGTH", 256)) - def aquila_attention_forward( self, @@ -75,58 +67,27 @@ def aquila_attention_forward( .transpose(1, 2) kv_seq_len = key_states.shape[-2] - enough_kv_room = True if past_key_value is not None: - enough_kv_room = is_enough_kv_cache_room_4_31(past_key_value, seq_len=kv_seq_len) kv_seq_len += past_key_value[0].shape[-2] - if query_states.device.type == "xpu" and not (self.training and query_states.requires_grad): - query_states, key_states = apply_rotary_pos_emb_no_cache_xpu(query_states, - key_states, - position_ids, - "aquila") + + if should_use_fuse_rope(hidden_states, position_ids, self.training): + import xe_addons + xe_addons.rotary_half_inplaced(self.rotary_emb.inv_freq, position_ids, + query_states, key_states) else: cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids, "aquila") - # [bsz, nh, t, hd] - if past_key_value is not None: - # reuse k, v, self_attention - cache_k = past_key_value[0] - cache_v = past_key_value[1] - if not enough_kv_room: - # allocate new - new_cache_k, new_cache_v = extend_kv_cache(bsz, - self.num_heads, # Support GQA - self.head_dim, - cache_k.size(2), - kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH, - dtype=cache_k.dtype, - device=hidden_states.device) - new_cache_k[:] = cache_k - new_cache_v[:] = cache_v - cache_k = new_cache_k - cache_v = new_cache_v - - key_states, value_states = append_kv_cache(cache_k, cache_v, key_states, value_states) - - elif use_cache: - max_cache_length = kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH - new_key_states, new_value_states = init_kv_cache(bsz, - self.num_heads, - self.head_dim, - kv_seq_len, - max_cache_length, - dtype=key_states.dtype, - device=hidden_states.device) - new_key_states[:] = key_states - new_value_states[:] = value_states - key_states = new_key_states - value_states = new_value_states + key_states, value_states = update_past_key_value( + past_key_value, key_states, value_states, + kv_seq_len, False, hidden_states.device + ) past_key_value = (key_states, value_states) if use_cache else None - attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + attn_weights = torch.matmul(query_states, + key_states.transpose(2, 3)) / math.sqrt(self.head_dim) attn_weights = torch.clamp(attn_weights, min=-1024., max=1024.) if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): @@ -148,8 +109,7 @@ def aquila_attention_forward( ) # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32)\ - .to(query_states.dtype) + attn_weights = attention_softmax(attn_weights, self.training) attn_output = torch.matmul(attn_weights, value_states) if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): diff --git a/python/llm/src/ipex_llm/transformers/models/decilm.py b/python/llm/src/ipex_llm/transformers/models/decilm.py index 99a8c8f4..194d735c 100644 --- a/python/llm/src/ipex_llm/transformers/models/decilm.py +++ b/python/llm/src/ipex_llm/transformers/models/decilm.py @@ -34,11 +34,11 @@ import torch from typing import Optional, Tuple import torch.nn.functional as F -from ipex_llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache -from ipex_llm.transformers.models.utils import is_enough_kv_cache_room_4_31, \ - apply_rotary_pos_emb +from ipex_llm.transformers.models.utils import apply_rotary_pos_emb from ipex_llm.transformers.models.utils import apply_rotary_pos_emb_no_cache_xpu -from ipex_llm.transformers.models.llama import should_use_fuse_rope, repeat_kv +from ipex_llm.transformers.models.llama import repeat_kv +from ipex_llm.transformers.models.utils import should_use_fuse_rope +from ipex_llm.transformers.models.utils import update_past_key_value from ipex_llm.utils.common import invalidInputError import os @@ -61,32 +61,9 @@ def decilm_attention_forward_4_35_2( is_decode = past_key_value is not None device = hidden_states.device - use_fuse_rope = should_use_fuse_rope(self, hidden_states, position_ids) - enough_kv_room = is_enough_kv_cache_room_4_31(past_key_value, seq_len=q_len) - - if self.config.pretraining_tp > 1: - key_value_slicing = ((self.num_key_value_heads * self.head_dim) // - self.config.pretraining_tp) - query_slices = self.q_proj.weight.split((self.num_heads * self.head_dim) - // self.config.pretraining_tp, dim=0) - key_slices = self.k_proj.weight.split(key_value_slicing, dim=0) - value_slices = self.v_proj.weight.split(key_value_slicing, dim=0) - - query_states = [F.linear(hidden_states, query_slices[i]) - for i in range(self.config.pretraining_tp)] - query_states = torch.cat(query_states, dim=-1) - - key_states = [F.linear(hidden_states, key_slices[i]) - for i in range(self.config.pretraining_tp)] - key_states = torch.cat(key_states, dim=-1) - - value_states = [F.linear(hidden_states, value_slices[i]) - for i in range(self.config.pretraining_tp)] - value_states = torch.cat(value_states, dim=-1) - else: - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) @@ -99,7 +76,7 @@ def decilm_attention_forward_4_35_2( if past_key_value is not None: kv_seq_len += past_key_value[0].shape[-2] - if use_fuse_rope: + if should_use_fuse_rope(hidden_states, position_ids, self.training): query_states, key_states = apply_rotary_pos_emb_no_cache_xpu(query_states, key_states, position_ids, @@ -109,39 +86,10 @@ def decilm_attention_forward_4_35_2( query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids, "llama") - if past_key_value is not None: - # reuse k, v, self_attention - cache_k = past_key_value[0] - cache_v = past_key_value[1] - if not enough_kv_room: - # allocate new - new_cache_k, new_cache_v = extend_kv_cache(bsz, - self.num_key_value_heads, # Support GQA - self.head_dim, - cache_k.size(2), - kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH, - dtype=cache_k.dtype, - device=device) - new_cache_k[:] = cache_k - new_cache_v[:] = cache_v - cache_k = new_cache_k - cache_v = new_cache_v - - key_states, value_states = append_kv_cache(cache_k, cache_v, key_states, value_states) - - elif use_cache: - max_cache_length = kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH - new_key_states, new_value_states = init_kv_cache(bsz, - self.num_key_value_heads, - self.head_dim, - kv_seq_len, - max_cache_length, - dtype=key_states.dtype, - device=device) - new_key_states[:] = key_states - new_value_states[:] = value_states - key_states = new_key_states - value_states = new_value_states + key_states, value_states = update_past_key_value( + past_key_value, key_states, value_states, + kv_seq_len, False, device + ) past_key_value = (key_states, value_states) if use_cache else None @@ -167,14 +115,8 @@ def decilm_attention_forward_4_35_2( attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size) - if self.config.pretraining_tp > 1: - attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2) - o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, - dim=1) - attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) - for i in range(self.config.pretraining_tp)]) - else: - attn_output = self.o_proj(attn_output) + attn_output = attn_output.to(hidden_states.dtype) + attn_output = self.o_proj(attn_output) if not output_attentions: attn_weights = None diff --git a/python/llm/src/ipex_llm/transformers/models/llama.py b/python/llm/src/ipex_llm/transformers/models/llama.py index 94a54349..4adbfcbd 100644 --- a/python/llm/src/ipex_llm/transformers/models/llama.py +++ b/python/llm/src/ipex_llm/transformers/models/llama.py @@ -46,16 +46,14 @@ from ipex_llm.transformers.models.utils import init_fp8_kv_cache, append_fp8_kv_ get_compresskv_attn_mask from ipex_llm.transformers.models.utils import is_enough_kv_cache_room_4_31, \ apply_rotary_pos_emb, is_enough_kv_cache_room_4_36 -from ipex_llm.transformers.models.utils import apply_rotary_pos_emb_no_cache_xpu from ipex_llm.transformers.models.utils import use_flash_attention, use_sdp, use_sdp_causal from ipex_llm.transformers.models.utils import mlp_fusion_check, fp16_fusion_check from ipex_llm.transformers.models.utils import use_decoding_fast_path, get_q_proj_or_qkv_proj from transformers.modeling_outputs import BaseModelOutputWithPast from transformers.models.llama.modeling_llama import LlamaModel, LlamaAttention from ipex_llm.transformers.low_bit_linear import SYM_INT4, FP8E5, IQ2_XXS, FP4 -from ipex_llm.ggml.quantize import ggml_tensor_qtype from ipex_llm.utils.common import invalidInputError -from ipex_llm.transformers.models.common import merge_qkv_base, fuse_mlp_base +from ipex_llm.transformers.models.common import merge_qkv_base try: from transformers.cache_utils import Cache, DynamicCache diff --git a/python/llm/src/ipex_llm/transformers/models/phixtral.py b/python/llm/src/ipex_llm/transformers/models/phixtral.py index d1029985..d1a72ca8 100644 --- a/python/llm/src/ipex_llm/transformers/models/phixtral.py +++ b/python/llm/src/ipex_llm/transformers/models/phixtral.py @@ -37,38 +37,8 @@ # limitations under the License. """ PyTorch Phixtral model.""" -import math -from typing import Optional, Tuple - import torch -from torch import nn import torch.nn.functional as F -from ipex_llm.ggml.quantize import ggml_tensor_qtype -from ipex_llm.utils.common import invalidInputError -from ipex_llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache -from ipex_llm.transformers.models.utils import apply_rotary_pos_emb,\ - apply_rotary_pos_emb_no_cache_xpu, is_enough_kv_cache_room_4_36 -from ipex_llm.transformers.models.mistral import should_use_fuse_rope -from ipex_llm.transformers.models.utils import use_flash_attention -from ipex_llm.transformers.models.utils import mlp_fusion_check - -import os - -KV_CACHE_ALLOC_BLOCK_LENGTH = int(os.environ.get("KV_CACHE_ALLOC_BLOCK_LENGTH", 256)) - - -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). - The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) - to (batch, num_attention_heads, seqlen, head_dim) - """ - batch, num_key_value_heads, slen, head_dim = hidden_states.shape - if n_rep == 1: - return hidden_states - hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, - n_rep, slen, head_dim) - return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) def phixtral_moeblock_forward(self, hidden_states: torch.Tensor):