From df2f3885ba8027b258cf1e8a0d2eaf5ef2aea514 Mon Sep 17 00:00:00 2001 From: SONG Ge <38711238+sgwhat@users.noreply.github.com> Date: Mon, 26 Feb 2024 11:29:48 +0800 Subject: [PATCH] [LLM] Enable kv_cache and forward_qkv optimizations for yuan2 (#10225) * add init kv_cache support for yuan2 * add forward qkv in yuan --- .../src/bigdl/llm/transformers/models/yuan.py | 149 ++++++++++-------- 1 file changed, 83 insertions(+), 66 deletions(-) diff --git a/python/llm/src/bigdl/llm/transformers/models/yuan.py b/python/llm/src/bigdl/llm/transformers/models/yuan.py index a48015b2..ed869a0c 100644 --- a/python/llm/src/bigdl/llm/transformers/models/yuan.py +++ b/python/llm/src/bigdl/llm/transformers/models/yuan.py @@ -20,13 +20,32 @@ # https://huggingface.co/IEITYuan/Yuan2-2B-hf/blob/7ab7b3c18eb8e5232ce2a3f720d4e6f4b53a2806/README.md#%E5%A3%B0%E6%98%8E%E4%B8%8E%E5%8D%8F%E8%AE%AEterms-and-conditions # -import torch import copy import math +from einops import rearrange from typing import Optional, Tuple -from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb +import torch +import torch.nn as nn from bigdl.llm.utils.common import invalidInputError +from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb +from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache +from bigdl.llm.transformers.models.utils import is_enough_kv_cache_room_4_31 +from bigdl.llm.transformers.low_bit_linear import SYM_INT4, FP8E5 + +KV_CACHE_ALLOC_BLOCK_LENGTH = 256 + + +def use_decoding_fast_path(q_type, use_fuse_rope, enough_kv_room, bs): + return q_type in [SYM_INT4, FP8E5] and \ + use_fuse_rope and enough_kv_room and bs == 1 + + +def should_use_fuse_rope(self, hidden_states, position_ids): + use_fuse_rope = hidden_states.device.type == "xpu" + use_fuse_rope = use_fuse_rope and not (self.training and hidden_states.requires_grad) + use_fuse_rope = use_fuse_rope and position_ids is not None + return use_fuse_rope def yuan_attention_forward( @@ -39,8 +58,13 @@ def yuan_attention_forward( use_cache: bool = False, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() + device = hidden_states.device before_hidden_states = None is_first_step = False + self.use_shareqk = False + + enough_kv_room = is_enough_kv_cache_room_4_31(past_key_value) + if use_cache: if past_key_value is None: inference_hidden_states_memory = torch.empty(bsz, 2, @@ -64,7 +88,9 @@ def yuan_attention_forward( value_states = \ self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + if self.use_shareqk: + # use_shareqk is disabled for now qk_states = self.qk_proj(hidden_states).view(bsz, q_len, self.num_heads*self.head_dim) query_key = qk_states.unsqueeze(2) * self.qk_weight + self.qk_bias query_states, key_states = torch.unbind(query_key, dim=2) @@ -95,80 +121,71 @@ def yuan_attention_forward( if past_key_value is not None: # reuse k, v, self_attention - key_states = torch.cat([past_key_value[0], key_states], dim=2) - value_states = torch.cat([past_key_value[1], value_states], dim=2) + 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, + self.head_dim, + cache_k.size(2), + kv_seq_len + KV_CACHE_ALLOC_BLOCK_LENGTH, + dtype=cache_k.d_type, + 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_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 past_key_value = \ (key_states, value_states, inference_hidden_states_memory) if use_cache else None - if self.use_flash_attention: - attn_weights = None - query_states = query_states.transpose(1, 2) - key_states = key_states.transpose(1, 2) - value_states = value_states.transpose(1, 2) + attn_weights = \ + torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) - batch_size, seqlen_q = query_states.shape[0], query_states.shape[1] - seqlen_k = key_states.shape[1] + invalidInputError(attn_weights.size() == (bsz, self.num_heads, q_len, kv_seq_len), + "Attention weights should be of size " + f"{(bsz, self.num_heads, q_len, kv_seq_len)}, " + f"but is {attn_weights.size()}") - q, k, v = \ - [rearrange(x, 'b s ... -> (b s) ...') for x in [query_states, key_states, value_states]] + if attention_mask is not None: + invalidInputError(attention_mask.size() == (bsz, 1, q_len, kv_seq_len), + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, " + f"but is {attention_mask.size()}") + attn_weights = attn_weights + attention_mask + attn_weights = torch.max(attn_weights, + torch.tensor(torch.finfo(attn_weights.dtype).min)) - cu_seqlens_q = torch.arange(0, (batch_size + 1) * seqlen_q, - step=seqlen_q, - dtype=torch.int, - device=q.device) + # upcast attention to fp32 + attn_weights = \ + torch.nn.functional.softmax(attn_weights, + dim=-1, + dtype=torch.float32).to(query_states.dtype) + attn_output = torch.matmul(attn_weights, value_states) - if self.training: - invalidInputError(seqlen_k == seqlen_q, - "`seqlen_k` should be equal to `seqlen_q`, but is not") - cu_seqlens_k = cu_seqlens_q - is_causal = self.causal_mask - else: - is_causal = seqlen_q == seqlen_k - cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, - step=seqlen_k, - dtype=torch.int, - device=q.device) - self.dropout = 0 - - output = flash_attn_unpadded_func( - q, k, v, cu_seqlens_q, cu_seqlens_k, seqlen_q, seqlen_k, self.dropout, causal=is_causal - ) - - attn_output = rearrange(output, '(b s) ... -> b s ...', b=batch_size) - else: - attn_weights = \ - torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) - - invalidInputError(attn_weights.size() == (bsz, self.num_heads, q_len, kv_seq_len), - "Attention weights should be of size " - f"{(bsz, self.num_heads, q_len, kv_seq_len)}, " - f"but is {attn_weights.size()}") - - if attention_mask is not None: - invalidInputError(attention_mask.size() == (bsz, 1, q_len, kv_seq_len), - f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, " - f"but is {attention_mask.size()}") - attn_weights = attn_weights + attention_mask - attn_weights = torch.max(attn_weights, - torch.tensor(torch.finfo(attn_weights.dtype).min)) - - # upcast attention to fp32 - attn_weights = \ - torch.nn.functional.softmax(attn_weights, - dim=-1, - dtype=torch.float32).to(query_states.dtype) - attn_output = torch.matmul(attn_weights, value_states) - - invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim), - "`attn_output` should be of size " - f"{(bsz, self.num_heads, q_len, self.head_dim)}, " - f"but is {attn_output.size()}") - - attn_output = attn_output.transpose(1, 2) + invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim), + "`attn_output` should be of size " + f"{(bsz, self.num_heads, q_len, self.head_dim)}, " + f"but is {attn_output.size()}") + attn_output = attn_output.transpose(1, 2) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) - attn_output = self.o_proj(attn_output) if not output_attentions: