679 lines
31 KiB
Python
679 lines
31 KiB
Python
#
|
||
# 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.
|
||
#
|
||
# This file is adapted from
|
||
# https://huggingface.co/THUDM/chatglm2-6b/blob/8eb45c842594b8473f291d0f94e7bbe86ffc67d8/modeling_chatglm.py
|
||
#
|
||
|
||
import math
|
||
import torch
|
||
from typing import Optional, Tuple, List
|
||
import torch.nn.functional as F
|
||
from transformers.modeling_outputs import BaseModelOutputWithPast
|
||
from ipex_llm.transformers.models.utils import init_kv_cache, extend_kv_cache, append_kv_cache
|
||
from ipex_llm.transformers.models.utils import init_fp8_kv_cache, append_fp8_kv_cache, \
|
||
restore_fp8_kv_cache, use_quantize_kv_cache
|
||
from ipex_llm.transformers.models.utils import use_sdp
|
||
|
||
|
||
import os
|
||
|
||
KV_CACHE_ALLOC_BLOCK_LENGTH = int(os.environ.get("KV_CACHE_ALLOC_BLOCK_LENGTH", 256))
|
||
KV_CACHE_ALLOC_MIN_LENGTH = 512
|
||
|
||
|
||
def split_tensor_along_last_dim(
|
||
tensor: torch.Tensor,
|
||
num_partitions: int,
|
||
contiguous_split_chunks: bool = False,
|
||
) -> List[torch.Tensor]:
|
||
"""Split a tensor along its last dimension.
|
||
Arguments:
|
||
tensor: input tensor.
|
||
num_partitions: number of partitions to split the tensor
|
||
contiguous_split_chunks: If True, make each chunk contiguous
|
||
in memory.
|
||
Returns:
|
||
A list of Tensors
|
||
"""
|
||
# Get the size and dimension.
|
||
last_dim = tensor.dim() - 1
|
||
last_dim_size = tensor.size()[last_dim] // num_partitions
|
||
# Split.
|
||
tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
|
||
# Note: torch.split does not create contiguous tensors by default.
|
||
if contiguous_split_chunks:
|
||
return tuple(chunk.contiguous() for chunk in tensor_list)
|
||
|
||
return tensor_list
|
||
|
||
|
||
@torch.jit.script
|
||
def apply_rotary_pos_emb_chatglm(x: torch.Tensor, rope_cache: torch.Tensor) -> torch.Tensor:
|
||
# x: [sq, b, np, hn]
|
||
sq, b, np, hn = x.size(0), x.size(1), x.size(2), x.size(3)
|
||
rot_dim = rope_cache.shape[-2] * 2
|
||
x, x_pass = x[..., :rot_dim], x[..., rot_dim:]
|
||
# truncate to support variable sizes
|
||
rope_cache = rope_cache[:sq]
|
||
xshaped = x.reshape(sq, -1, np, rot_dim // 2, 2)
|
||
rope_cache = rope_cache.view(sq, -1, 1, xshaped.size(3), 2)
|
||
x_out2 = torch.stack(
|
||
[
|
||
xshaped[..., 0] * rope_cache[..., 0] - xshaped[..., 1] * rope_cache[..., 1],
|
||
xshaped[..., 1] * rope_cache[..., 0] + xshaped[..., 0] * rope_cache[..., 1],
|
||
],
|
||
-1,
|
||
)
|
||
x_out2 = x_out2.flatten(3)
|
||
return torch.cat((x_out2, x_pass), dim=-1)
|
||
|
||
|
||
def repeat_kv(key: torch.Tensor, value: torch.Tensor, n_head: int) -> (torch.Tensor, torch.Tensor):
|
||
# key, value's shape: [bs, n_kv_head, seq_len, head_dim] -> [bs, n_head, seq_len, head_dim]
|
||
batch_size, n_kv_head, seq_len, head_dim = key.shape
|
||
|
||
key = key.unsqueeze(2)
|
||
key = key.expand(-1, -1, n_head // n_kv_head, -1, -1)
|
||
key = key.contiguous().view(batch_size, n_head, seq_len, head_dim)
|
||
|
||
value = value.unsqueeze(2)
|
||
value = value.expand(-1, -1, n_head // n_kv_head, -1, -1)
|
||
value = value.contiguous().view(batch_size, n_head, seq_len, head_dim)
|
||
|
||
return key, value
|
||
|
||
|
||
def should_split_qkv_tensor(query_layer, bsz, n_head, seq_len):
|
||
if os.environ.get("IPEX_LLM_SPLIT_QKV", None) is not None:
|
||
return os.environ.get("IPEX_LLM_SPLIT_QKV", None) == "1"
|
||
elif query_layer.dtype == torch.float16 and query_layer.shape[2] >= 5000:
|
||
# split tensor for memory block limitation
|
||
# support fp16 and set input length threshold at 5000 for now
|
||
return True
|
||
elif query_layer.element_size()*bsz*n_head*seq_len*seq_len >= 4*1024**3:
|
||
# attn_weight size larger than memory block limitation 4GB
|
||
return True
|
||
return False
|
||
|
||
|
||
def chatglm_rms_norm_forward(self, hidden_states):
|
||
if hidden_states.device.type == "xpu" and not (self.training and hidden_states.requires_grad):
|
||
import linear_q4_0
|
||
x_2d = hidden_states.reshape(-1, hidden_states.size(-1)).contiguous()
|
||
output = linear_q4_0.rms_norm(self.weight, x_2d, self.eps)
|
||
return output.reshape(hidden_states.shape)
|
||
|
||
input_dtype = hidden_states.dtype
|
||
hidden_states = hidden_states.to(torch.float32)
|
||
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
||
hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
|
||
return self.weight * hidden_states.to(input_dtype)
|
||
|
||
|
||
def chatglm2_model_forward(
|
||
self,
|
||
input_ids,
|
||
position_ids: Optional[torch.Tensor]=None,
|
||
attention_mask: Optional[torch.BoolTensor]=None,
|
||
full_attention_mask: Optional[torch.BoolTensor]=None,
|
||
past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]]=None,
|
||
inputs_embeds: Optional[torch.Tensor]=None,
|
||
use_cache: Optional[bool]=None,
|
||
output_hidden_states: Optional[bool]=None,
|
||
return_dict: Optional[bool]=None,
|
||
):
|
||
output_hidden_states = (
|
||
output_hidden_states if output_hidden_states is not None
|
||
else self.config.output_hidden_states
|
||
)
|
||
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||
|
||
batch_size, seq_length = input_ids.shape
|
||
|
||
if inputs_embeds is None:
|
||
inputs_embeds = self.embedding(input_ids)
|
||
|
||
if full_attention_mask is None:
|
||
if (attention_mask is not None and not attention_mask.all()) or (
|
||
past_key_values and seq_length != 1):
|
||
full_attention_mask = self.get_masks(input_ids,
|
||
past_key_values,
|
||
padding_mask=attention_mask)
|
||
|
||
use_fuse_rope = input_ids.device.type == "xpu"
|
||
use_fuse_rope = use_fuse_rope and not self.training
|
||
|
||
# Rotary positional embeddings
|
||
rotary_pos_emb = self.rotary_pos_emb(self.seq_length)
|
||
if position_ids is not None:
|
||
rotary_pos_emb = rotary_pos_emb[position_ids]
|
||
else:
|
||
rotary_pos_emb = rotary_pos_emb[None, :seq_length]
|
||
if use_fuse_rope:
|
||
# Repeat cos sin here, call only once for each token.
|
||
# Chatglm2's rotary embedding is similar to gptj's, is rotate_every_two.
|
||
# If put this to attension forward, it will generate too many times.
|
||
cos, sin = rotary_pos_emb.split(rotary_pos_emb.shape[-1] // 2, dim=-1)
|
||
cos = cos.squeeze(-1)
|
||
sin = sin.squeeze(-1)
|
||
cos = torch.repeat_interleave(cos[:, :, None, :], 2, 3)
|
||
sin = torch.repeat_interleave(sin[:, :, None, :], 2, 3)
|
||
rotary_pos_emb = (cos, sin)
|
||
else:
|
||
rotary_pos_emb = rotary_pos_emb.transpose(0, 1).contiguous()
|
||
|
||
# Run encoder.
|
||
hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
|
||
inputs_embeds, full_attention_mask, rotary_pos_emb=rotary_pos_emb,
|
||
kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
|
||
)
|
||
|
||
if not return_dict:
|
||
return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions]
|
||
if v is not None)
|
||
|
||
return BaseModelOutputWithPast(
|
||
last_hidden_state=hidden_states,
|
||
past_key_values=presents,
|
||
hidden_states=all_hidden_states,
|
||
attentions=all_self_attentions,
|
||
)
|
||
|
||
|
||
def chatglm2_attention_forward(
|
||
self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
|
||
):
|
||
if use_quantize_kv_cache(self.query_key_value, hidden_states.transpose(0, 1)):
|
||
forward_function = chatglm2_quantized_attention_forward_8eb45c
|
||
else:
|
||
forward_function = chatglm2_attention_forward_8eb45c
|
||
return forward_function(
|
||
self=self,
|
||
hidden_states=hidden_states,
|
||
attention_mask=attention_mask,
|
||
rotary_pos_emb=rotary_pos_emb,
|
||
kv_cache=kv_cache,
|
||
use_cache=use_cache
|
||
)
|
||
|
||
|
||
def chatglm2_quantized_attention_forward_8eb45c(
|
||
self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
|
||
):
|
||
# hidden_states: [seq_len, bs, head_dim]
|
||
mixed_x_layer = self.query_key_value(hidden_states)
|
||
|
||
n_head = self.num_attention_heads_per_partition
|
||
n_kv_head = self.num_multi_query_groups_per_partition if self.multi_query_attention else n_head
|
||
head_dim = self.hidden_size_per_attention_head
|
||
|
||
query_layer, key_layer, value_layer = mixed_x_layer.split(
|
||
[n_head * head_dim, n_kv_head * head_dim, n_kv_head * head_dim],
|
||
dim=-1,
|
||
)
|
||
query_layer = query_layer.view(query_layer.shape[:-1] + (n_head, head_dim))
|
||
key_layer = key_layer.view(key_layer.shape[:-1] + (n_kv_head, head_dim))
|
||
value_layer = value_layer.view(value_layer.shape[:-1] + (n_kv_head, head_dim))
|
||
# query, key, value's shape: [seq_len, bs, n_head/n_kv_head, head_dim]
|
||
|
||
# apply relative positional encoding (rotary embedding)
|
||
if rotary_pos_emb is not None:
|
||
if len(rotary_pos_emb) == 2 and isinstance(rotary_pos_emb, tuple):
|
||
# use_fuse_rope, see chatglm2_model_forward
|
||
cos, sin = rotary_pos_emb
|
||
rot_dim = cos.shape[-1]
|
||
query_layer = query_layer.transpose(0, 1)
|
||
key_layer = key_layer.transpose(0, 1)
|
||
query_layer_cur = query_layer[..., :rot_dim]
|
||
key_layer_cur = key_layer[..., :rot_dim]
|
||
# ipex_llm's apply_rotary_embedding can change the origin storage,
|
||
# so query_layer will get the result directly.
|
||
torch.ops.torch_ipex.apply_rotary_embedding(query_layer_cur, sin, cos, query_layer_cur)
|
||
torch.ops.torch_ipex.apply_rotary_embedding(key_layer_cur, sin, cos, key_layer_cur)
|
||
query_layer = query_layer.transpose(0, 1)
|
||
key_layer = key_layer.transpose(0, 1)
|
||
else:
|
||
query_layer = apply_rotary_pos_emb_chatglm(query_layer, rotary_pos_emb)
|
||
key_layer = apply_rotary_pos_emb_chatglm(key_layer, rotary_pos_emb)
|
||
|
||
query_layer = query_layer.permute(1, 2, 0, 3)
|
||
key_layer = key_layer.permute(1, 2, 0, 3)
|
||
value_layer = value_layer.permute(1, 2, 0, 3)
|
||
# query, key, value's shape: [bs, n_head/n_kv_head, seq_len, head_dim]
|
||
batch_size, _, seq_len, _ = query_layer.shape
|
||
|
||
if kv_cache is None:
|
||
# first token
|
||
if self.multi_query_attention:
|
||
key, value = repeat_kv(key_layer, value_layer, n_head)
|
||
else:
|
||
key, value = key_layer, value_layer
|
||
|
||
if should_split_qkv_tensor(query_layer, batch_size, n_head, seq_len):
|
||
# split second dim to block size = 8
|
||
block_size = 8
|
||
query_split = torch.split(query_layer, block_size, dim=1)
|
||
key_split = torch.split(key, block_size, dim=1)
|
||
value_split = torch.split(value, block_size, dim=1)
|
||
results = []
|
||
for q, k, v in zip(query_split, key_split, value_split):
|
||
if attention_mask is None:
|
||
result = F.scaled_dot_product_attention(q, k, v, is_causal=True)
|
||
else:
|
||
result = F.scaled_dot_product_attention(q, k, v, attention_mask)
|
||
results.append(result)
|
||
context_layer = torch.cat(results, dim=1)
|
||
else:
|
||
if attention_mask is None:
|
||
context_layer = F.scaled_dot_product_attention(query_layer, key,
|
||
value, is_causal=True)
|
||
else:
|
||
context_layer = F.scaled_dot_product_attention(query_layer, key,
|
||
value, attention_mask)
|
||
context_layer = context_layer.to(query_layer.dtype)
|
||
|
||
if use_cache:
|
||
k_cache, v_cache = init_fp8_kv_cache(batch_size,
|
||
n_kv_head,
|
||
seq_len,
|
||
head_dim,
|
||
query_layer.device)
|
||
k_cache, v_cache = append_fp8_kv_cache(k_cache, v_cache, key_layer, value_layer)
|
||
else:
|
||
k_cache, v_cache = kv_cache
|
||
k_cache = k_cache.permute(1, 2, 0, 3)
|
||
v_cache = v_cache.permute(1, 2, 0, 3)
|
||
# k_cache, v_cache's shape: [bs, n_kv_head, seq_len, head_dim]
|
||
|
||
k_cache, v_cache = append_fp8_kv_cache(k_cache, v_cache, key_layer, value_layer)
|
||
|
||
if attention_mask is not None:
|
||
attention_mask = ~attention_mask
|
||
attn_bias = torch.zeros(attention_mask.shape, dtype=query_layer.dtype,
|
||
device=query_layer.device)
|
||
if attention_mask.dtype == torch.bool:
|
||
attn_bias.masked_fill_(attention_mask.logical_not(), float("-inf"))
|
||
else:
|
||
attn_bias += attention_mask
|
||
else:
|
||
attn_bias = None
|
||
|
||
if seq_len != 1:
|
||
key, value = restore_fp8_kv_cache(k_cache, v_cache, query_layer.dtype)
|
||
key, value = repeat_kv(key, value, n_head)
|
||
attn = torch.matmul(query_layer, key.transpose(2, 3)) / math.sqrt(head_dim)
|
||
if attn_bias is not None:
|
||
attn += attn_bias
|
||
attn = F.softmax(attn, dim=-1, dtype=torch.float32)
|
||
context_layer = torch.matmul(attn.to(value.dtype), value)
|
||
else:
|
||
key, value = k_cache, v_cache
|
||
import linear_q4_0
|
||
context_layer = linear_q4_0.sdp_fp8(query_layer, key, value, attn_bias)
|
||
|
||
# context_layer's shape: [bs, n_head, seq_len, head_dim] -> [seq_len, bs, n_head * head_dim]
|
||
context_layer = context_layer.permute(2, 0, 1, 3).contiguous().view(seq_len, batch_size, -1)
|
||
|
||
if use_cache:
|
||
kv_cache = (k_cache.permute(2, 0, 1, 3), v_cache.permute(2, 0, 1, 3))
|
||
else:
|
||
kv_cache = None
|
||
|
||
output = self.dense(context_layer)
|
||
|
||
return output, kv_cache
|
||
|
||
|
||
def chatglm2_attention_forward_8eb45c(
|
||
self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
|
||
):
|
||
# hidden_states: [sq, b, h]
|
||
|
||
# =================================================
|
||
# Pre-allocate memory for key-values for inference.
|
||
# =================================================
|
||
# =====================
|
||
# Query, Key, and Value
|
||
# =====================
|
||
|
||
# Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
|
||
device = hidden_states.device
|
||
mixed_x_layer = self.query_key_value(hidden_states)
|
||
|
||
if self.multi_query_attention:
|
||
(query_layer, key_layer, value_layer) = mixed_x_layer.split(
|
||
[
|
||
self.num_attention_heads_per_partition * self.hidden_size_per_attention_head,
|
||
self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
|
||
self.num_multi_query_groups_per_partition * self.hidden_size_per_attention_head,
|
||
],
|
||
dim=-1,
|
||
)
|
||
query_layer = query_layer.view(
|
||
query_layer.size()[:-1] + (self.num_attention_heads_per_partition,
|
||
self.hidden_size_per_attention_head)
|
||
)
|
||
key_layer = key_layer.view(
|
||
key_layer.size()[:-1] + (self.num_multi_query_groups_per_partition,
|
||
self.hidden_size_per_attention_head)
|
||
)
|
||
value_layer = value_layer.view(
|
||
value_layer.size()[:-1]
|
||
+ (self.num_multi_query_groups_per_partition, self.hidden_size_per_attention_head)
|
||
)
|
||
else:
|
||
new_tensor_shape = mixed_x_layer.size()[:-1] + (self.num_attention_heads_per_partition,
|
||
3 * self.hidden_size_per_attention_head)
|
||
mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
|
||
|
||
# [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
|
||
(query_layer, key_layer, value_layer) = split_tensor_along_last_dim(mixed_x_layer, 3)
|
||
|
||
cur_length, batch_size = query_layer.shape[0], query_layer.shape[1]
|
||
|
||
# apply relative positional encoding (rotary embedding)
|
||
if rotary_pos_emb is not None:
|
||
if len(rotary_pos_emb) == 2 and isinstance(rotary_pos_emb, tuple):
|
||
# use_fuse_rope, see chatglm2_model_forward
|
||
cos, sin = rotary_pos_emb
|
||
rot_dim = cos.shape[-1]
|
||
query_layer = query_layer.transpose(0, 1)
|
||
key_layer = key_layer.transpose(0, 1)
|
||
query_layer_cur = query_layer[..., :rot_dim]
|
||
key_layer_cur = key_layer[..., :rot_dim]
|
||
# ipex_llm's apply_rotary_embedding can change the origin storage,
|
||
# so query_layer will get the result directly.
|
||
torch.ops.torch_ipex.apply_rotary_embedding(query_layer_cur, sin, cos, query_layer_cur)
|
||
torch.ops.torch_ipex.apply_rotary_embedding(key_layer_cur, sin, cos, key_layer_cur)
|
||
query_layer = query_layer.transpose(0, 1)
|
||
key_layer = key_layer.transpose(0, 1)
|
||
else:
|
||
query_layer = apply_rotary_pos_emb_chatglm(query_layer, rotary_pos_emb)
|
||
key_layer = apply_rotary_pos_emb_chatglm(key_layer, rotary_pos_emb)
|
||
|
||
if self.multi_query_attention:
|
||
if device.type == "xpu" and batch_size > 1: # use beam_search for generation.
|
||
# If batch_size > 1 on gpu, permute key/value_layer to [bs, np, sl, hn]
|
||
# to reduce memory usage. Otherwise，expend key/value_layer to [bs, nh, sl, hn].
|
||
key_layer = key_layer.permute(1, 2, 0, 3) # [bs, np, sl, hn]
|
||
value_layer = value_layer.permute(1, 2, 0, 3) # [bs, np, sl, hn]
|
||
else:
|
||
key_length = key_layer.size(0)
|
||
query_group_size = self.num_attention_heads_per_partition // \
|
||
self.num_multi_query_groups_per_partition
|
||
key_layer = key_layer.permute(1, 2, 0, 3).unsqueeze(-3) # [bs, nh/k, sl, hn]
|
||
key_layer = key_layer.expand(-1, -1, query_group_size, -1, -1)
|
||
key_layer = key_layer.contiguous().view((batch_size,
|
||
self.num_attention_heads_per_partition,
|
||
key_length,
|
||
self.hidden_size_per_attention_head))
|
||
value_layer = value_layer.permute(1, 2, 0, 3).unsqueeze(-3) # [bs, nh/k, sl, hn]
|
||
value_layer = value_layer.expand(-1, -1, query_group_size, -1, -1)
|
||
value_layer = value_layer.contiguous().view((batch_size,
|
||
self.num_attention_heads_per_partition,
|
||
key_length,
|
||
self.hidden_size_per_attention_head))
|
||
|
||
# adjust key and value for inference
|
||
if kv_cache is not None:
|
||
cache_k, cache_v = kv_cache
|
||
cache_k = cache_k.permute(1, 2, 0, 3)
|
||
cache_v = cache_v.permute(1, 2, 0, 3)
|
||
past_length = cache_k.size(2)
|
||
|
||
if cache_k.stride()[1] < (past_length + cur_length) * cache_k.size(3):
|
||
max_cache_length = past_length + cur_length + KV_CACHE_ALLOC_BLOCK_LENGTH
|
||
if device.type == "xpu" and batch_size > 1: # use beam_search for generation.
|
||
# If batch_size > 1 on gpu, use init_kv_cache to avoid empty cache for ensuring
|
||
# generation correctness.
|
||
# Set the num_heads in init_kv_cache to np, ensuring that the tensors of
|
||
# new_cache_k/v and key/value_layer have the same size.
|
||
new_cache_k, new_cache_v = init_kv_cache(batch_size,
|
||
self.num_multi_query_groups_per_partition,
|
||
self.hidden_size_per_attention_head,
|
||
past_length,
|
||
max_cache_length,
|
||
dtype=query_layer.dtype,
|
||
device=device)
|
||
else:
|
||
new_cache_k, new_cache_v = extend_kv_cache(batch_size,
|
||
self.num_attention_heads_per_partition,
|
||
self.hidden_size_per_attention_head,
|
||
past_length,
|
||
max_cache_length,
|
||
dtype=query_layer.dtype,
|
||
device=device)
|
||
new_cache_k[:] = cache_k
|
||
new_cache_v[:] = cache_v
|
||
cache_k = new_cache_k
|
||
cache_v = new_cache_v
|
||
|
||
key_layer, value_layer = append_kv_cache(cache_k, cache_v, key_layer, value_layer)
|
||
|
||
elif use_cache:
|
||
max_cache_length = max(KV_CACHE_ALLOC_MIN_LENGTH, cur_length) \
|
||
+ KV_CACHE_ALLOC_BLOCK_LENGTH
|
||
|
||
if device.type == "xpu" and batch_size > 1: # use beam_search for generation.
|
||
# Ensure the tensors of key/value_cache and key/value_layer have the same size.
|
||
nums_per_partition = self.num_multi_query_groups_per_partition
|
||
else:
|
||
nums_per_partition = self.num_attention_heads_per_partition
|
||
|
||
key_cache, value_cache = init_kv_cache(batch_size,
|
||
nums_per_partition,
|
||
self.hidden_size_per_attention_head,
|
||
cur_length,
|
||
max_cache_length,
|
||
dtype=query_layer.dtype,
|
||
device=device)
|
||
key_cache[:] = key_layer
|
||
value_cache[:] = value_layer
|
||
key_layer = key_cache
|
||
value_layer = value_cache
|
||
|
||
# If batch_size > 1, return tensors with shape [bs, np, sl, hn] as past_key_values. This could
|
||
# reduce memory usage as tensors are not expended to [bs, nh, sl, hn].
|
||
# Otherwise, return views of [bs, nh, sl, hn].
|
||
cache_key_layer = key_layer
|
||
cache_value_layer = value_layer
|
||
|
||
if use_cache:
|
||
kv_cache = (key_layer, value_layer)
|
||
else:
|
||
kv_cache = None
|
||
|
||
# ==================================
|
||
# core attention computation
|
||
# ==================================
|
||
if device.type == "xpu" and batch_size > 1: # use beam_search for generation.
|
||
# If batch_size > 1, expend key/value_layer to [ns, nh, sl, bn] for
|
||
# core attention computation.
|
||
# The expanded tensors will not be returned as past_key_values.
|
||
if self.multi_query_attention:
|
||
query_group_size = self.num_attention_heads_per_partition // \
|
||
self.num_multi_query_groups_per_partition
|
||
key_layer = key_layer.unsqueeze(-3)
|
||
key_layer = key_layer.expand(-1, -1, query_group_size, -1, -1)
|
||
save_length = key_layer.size(3)
|
||
# [bs, np, sl, hn] --> [bs, nh, sl, hn]
|
||
key_layer = key_layer.contiguous().view((batch_size,
|
||
self.num_attention_heads_per_partition,
|
||
save_length,
|
||
self.hidden_size_per_attention_head))
|
||
value_layer = value_layer.unsqueeze(-3)
|
||
value_layer = value_layer.expand(-1, -1, query_group_size, -1, -1)
|
||
# [bs, np, sl, hn] --> [bs, nh, sl, hn]
|
||
value_layer = value_layer.contiguous().view((batch_size,
|
||
self.num_attention_heads_per_partition,
|
||
save_length,
|
||
self.hidden_size_per_attention_head))
|
||
|
||
context_layer = core_attn_forward_8eb45c(query_layer, key_layer, value_layer, attention_mask)
|
||
|
||
# =================
|
||
# Output. [sq, b, h]
|
||
# =================
|
||
|
||
output = self.dense(context_layer)
|
||
|
||
return output, (cache_key_layer.permute(2, 0, 1, 3), cache_value_layer.permute(2, 0, 1, 3))
|
||
|
||
|
||
def core_attn_forward_8eb45c(query_layer, key_layer, value_layer, attention_mask):
|
||
pytorch_major_version = int(torch.__version__.split('.')[0])
|
||
if pytorch_major_version >= 2:
|
||
query_layer = query_layer.permute(1, 2, 0, 3)
|
||
L, S = query_layer.shape[2], key_layer.shape[2]
|
||
if attention_mask is None and L == S:
|
||
batch_size, n_head, seq_len, head_dim = query_layer.shape
|
||
if should_split_qkv_tensor(query_layer, batch_size, n_head, seq_len):
|
||
# split second dim to block size = 8
|
||
block_size = 8
|
||
query_split = torch.split(query_layer.to(key_layer.dtype), block_size, dim=1)
|
||
key_split = torch.split(key_layer, block_size, dim=1)
|
||
value_split = torch.split(value_layer, block_size, dim=1)
|
||
results = []
|
||
for q, k, v in zip(query_split, key_split, value_split):
|
||
result = F.scaled_dot_product_attention(q, k, v, is_causal=True).to(k.dtype)
|
||
results.append(result)
|
||
context_layer = torch.cat(results, dim=1)
|
||
else:
|
||
context_layer = F.scaled_dot_product_attention(query_layer.to(key_layer.dtype),
|
||
key_layer,
|
||
value_layer,
|
||
is_causal=True).to(key_layer.dtype)
|
||
else:
|
||
# attention_mask is not None only when past_key_value is not None and q_len > 1
|
||
if attention_mask is not None:
|
||
attn_bias = torch.zeros(attention_mask.shape, dtype=query_layer.dtype,
|
||
device=query_layer.device)
|
||
attention_mask = ~attention_mask
|
||
if attention_mask.dtype == torch.bool:
|
||
attn_bias.masked_fill_(attention_mask.logical_not(), float("-inf"))
|
||
else:
|
||
attn_bias += attention_mask
|
||
else:
|
||
attn_bias = None
|
||
|
||
if use_sdp(query_layer.shape[2], key_layer.shape[2],
|
||
query_layer.shape[-1], query_layer):
|
||
import linear_q4_0
|
||
attn_output = linear_q4_0.sdp(query_layer, key_layer, value_layer, attn_bias)
|
||
context_layer = attn_output.view(query_layer.shape)
|
||
else:
|
||
head_dim = query_layer.size(-1)
|
||
attn = torch.matmul(query_layer.to(key_layer.dtype),
|
||
key_layer.transpose(2, 3)) / math.sqrt(head_dim)
|
||
if attn_bias is not None:
|
||
attn += attn_bias
|
||
attn = F.softmax(attn, dim=-1,
|
||
dtype=torch.float32).to(value_layer.dtype)
|
||
context_layer = torch.matmul(attn, value_layer)
|
||
context_layer = context_layer.permute(2, 0, 1, 3)
|
||
new_context_layer_shape = context_layer.size()[:-2] + (-1,)
|
||
context_layer = context_layer.reshape(*new_context_layer_shape)
|
||
else:
|
||
# Raw attention scores
|
||
|
||
# [b, np, sq, sk]
|
||
output_size = (query_layer.size(1), query_layer.size(2),
|
||
query_layer.size(0), key_layer.size(2))
|
||
|
||
# [sq, b, np, hn] -> [sq, b * np, hn]
|
||
query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
|
||
# [sk, b, np, hn] -> [sk, b * np, hn]
|
||
key_layer = key_layer.view(output_size[0] * output_size[1], output_size[3], -1)
|
||
|
||
# preallocting input tensor: [b * np, sq, sk]
|
||
matmul_input_buffer = torch.empty(
|
||
output_size[0] * output_size[1],
|
||
output_size[2], output_size[3], dtype=query_layer.dtype,
|
||
device=query_layer.device
|
||
)
|
||
|
||
matmul_result = torch.empty(
|
||
output_size[0] * output_size[1],
|
||
output_size[2], output_size[3], dtype=query_layer.dtype,
|
||
device=query_layer.device
|
||
)
|
||
|
||
# Raw attention scores. [b * np, sq, sk]
|
||
torch.baddbmm(
|
||
matmul_input_buffer,
|
||
query_layer.transpose(0, 1), # [b * np, sq, hn]
|
||
key_layer.transpose(1, 2), # [b * np, hn, sk]
|
||
beta=0.0,
|
||
alpha=(1.0 / self.norm_factor),
|
||
out=matmul_result
|
||
)
|
||
|
||
# change view to [b, np, sq, sk]
|
||
attention_scores = matmul_result.view(*output_size)
|
||
|
||
# ===========================
|
||
# Attention probs and dropout
|
||
# ===========================
|
||
|
||
# attention scores and attention mask [b, np, sq, sk]
|
||
if self.attention_softmax_in_fp32:
|
||
attention_scores = attention_scores.float()
|
||
if self.coeff is not None:
|
||
attention_scores = attention_scores * self.coeff
|
||
if attention_mask is None and attention_scores.shape[2] == attention_scores.shape[3]:
|
||
attention_mask = torch.ones(output_size[0], 1, output_size[2], output_size[3],
|
||
device=attention_scores.device, dtype=torch.bool)
|
||
attention_mask.tril_()
|
||
attention_mask = ~attention_mask
|
||
if attention_mask is not None:
|
||
attention_scores = attention_scores.masked_fill(attention_mask, float("-inf"))
|
||
attention_probs = F.softmax(attention_scores, dim=-1)
|
||
attention_probs = attention_probs.type_as(value_layer)
|
||
|
||
# This is actually dropping out entire tokens to attend to, which might
|
||
# seem a bit unusual, but is taken from the original Transformer paper.
|
||
attention_probs = self.attention_dropout(attention_probs)
|
||
# =========================
|
||
# Context layer. [sq, b, hp]
|
||
# =========================
|
||
|
||
# value_layer -> context layer.
|
||
# [sk, b, np, hn] --> [b, np, sq, hn]
|
||
|
||
# context layer shape: [b, np, sq, hn]
|
||
output_size = (value_layer.size(0), value_layer.size(1),
|
||
query_layer.size(0), value_layer.size(3))
|
||
# change view [sk, b * np, hn]
|
||
value_layer = value_layer.view(output_size[0] * output_size[1], value_layer.size(2), -1)
|
||
# change view [b * np, sq, sk]
|
||
attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
|
||
# matmul: [b * np, sq, hn]
|
||
context_layer = torch.empty(
|
||
output_size[0] * output_size[1],
|
||
output_size[2], value_layer.size(-1), dtype=value_layer.dtype,
|
||
device=value_layer.device,
|
||
)
|
||
torch.bmm(attention_probs, value_layer, out=context_layer)
|
||
# change view [b, np, sq, hn]
|
||
context_layer = context_layer.view(*output_size)
|
||
# [b, np, sq, hn] --> [sq, b, np, hn]
|
||
context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
|
||
# [sq, b, np, hn] --> [sq, b, hp]
|
||
new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
|
||
context_layer = context_layer.view(*new_context_layer_shape)
|
||
|
||
return context_layer
|