fix and optimize chatglm2-32k and chatglm3-128k (#11306)
This commit is contained in:
Yishuo Wang
GitHub
parent
60cb1dac7c
commit
5e25766855
3 changed files with 57 additions and 217 deletions
|
|
@ -1008,26 +1008,21 @@ def _optimize_post(model, lightweight_bmm=False):
|
||||||
|
|
||||||
if model.config.architectures is not None \
|
if model.config.architectures is not None \
|
||||||
and model.config.architectures[0] in ["ChatGLMModel", "ChatGLMForConditionalGeneration"]:
|
and model.config.architectures[0] in ["ChatGLMModel", "ChatGLMForConditionalGeneration"]:
|
||||||
if (model.config.num_layers == 28 and hasattr(model.config, 'rope_ratio')
|
if hasattr(model.config, 'padded_vocab_size') and \
|
||||||
and model.config.rope_ratio == 16):
|
|
||||||
# chatglm2-6b-32k
|
|
||||||
modeling_module_name = model.__class__.__module__
|
|
||||||
module = importlib.import_module(modeling_module_name)
|
|
||||||
from ipex_llm.transformers.models.chatglm2_32k import chatglm2_32k_attention_forward
|
|
||||||
convert_forward(model,
|
|
||||||
module.SelfAttention,
|
|
||||||
chatglm2_32k_attention_forward)
|
|
||||||
elif hasattr(model.config, 'padded_vocab_size') and \
|
|
||||||
model.config.padded_vocab_size == 65024:
|
model.config.padded_vocab_size == 65024:
|
||||||
# chatglm2-6b
|
# chatglm2-6b, chatglm2-6b-32k, chatglm3-6b, chatglm3-6b-32k, chatglm3-6b-128k
|
||||||
modeling_module_name = model.__class__.__module__
|
modeling_module_name = model.__class__.__module__
|
||||||
module = importlib.import_module(modeling_module_name)
|
module = importlib.import_module(modeling_module_name)
|
||||||
from ipex_llm.transformers.models.chatglm2 import chatglm2_attention_forward
|
from ipex_llm.transformers.models.chatglm2 import chatglm2_attention_forward
|
||||||
from ipex_llm.transformers.models.chatglm2 import chatglm_rms_norm_forward
|
from ipex_llm.transformers.models.chatglm2 import chatglm_rms_norm_forward
|
||||||
|
from ipex_llm.transformers.models.chatglm2 import chatglm2_encoder_forward
|
||||||
from ipex_llm.transformers.models.chatglm2 import chatglm2_model_forward
|
from ipex_llm.transformers.models.chatglm2 import chatglm2_model_forward
|
||||||
convert_forward(model,
|
convert_forward(model,
|
||||||
module.SelfAttention,
|
module.SelfAttention,
|
||||||
chatglm2_attention_forward)
|
chatglm2_attention_forward)
|
||||||
|
convert_forward(model,
|
||||||
|
module.GLMTransformer,
|
||||||
|
chatglm2_encoder_forward)
|
||||||
convert_forward(model,
|
convert_forward(model,
|
||||||
module.ChatGLMModel,
|
module.ChatGLMModel,
|
||||||
chatglm2_model_forward)
|
chatglm2_model_forward)
|
||||||
|
|
|
||||||
|
|
@ -145,6 +145,57 @@ def chatglm2_model_forward(
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# remove code which stores first token's kv cache by tensor format
|
||||||
|
# to fix chatglm2-32k and chatglm3-128k
|
||||||
|
def chatglm2_encoder_forward(
|
||||||
|
self, hidden_states, attention_mask, rotary_pos_emb, kv_caches=None,
|
||||||
|
use_cache: Optional[bool] = True,
|
||||||
|
output_hidden_states: Optional[bool] = False,
|
||||||
|
):
|
||||||
|
if not kv_caches:
|
||||||
|
kv_caches = [None for _ in range(self.num_layers)]
|
||||||
|
presents = () if use_cache else None
|
||||||
|
if self.gradient_checkpointing and self.training:
|
||||||
|
use_cache = False
|
||||||
|
|
||||||
|
all_self_attentions = None
|
||||||
|
all_hidden_states = () if output_hidden_states else None
|
||||||
|
for index in range(self.num_layers):
|
||||||
|
if output_hidden_states:
|
||||||
|
all_hidden_states = all_hidden_states + (hidden_states,)
|
||||||
|
|
||||||
|
layer = self._get_layer(index)
|
||||||
|
if self.gradient_checkpointing and self.training:
|
||||||
|
layer_ret = torch.utils.checkpoint.checkpoint(
|
||||||
|
layer,
|
||||||
|
hidden_states,
|
||||||
|
attention_mask,
|
||||||
|
rotary_pos_emb,
|
||||||
|
kv_caches[index],
|
||||||
|
use_cache
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
layer_ret = layer(
|
||||||
|
hidden_states,
|
||||||
|
attention_mask,
|
||||||
|
rotary_pos_emb,
|
||||||
|
kv_cache=kv_caches[index],
|
||||||
|
use_cache=use_cache
|
||||||
|
)
|
||||||
|
hidden_states, kv_cache = layer_ret
|
||||||
|
if use_cache:
|
||||||
|
presents = presents + (kv_cache,)
|
||||||
|
|
||||||
|
if output_hidden_states:
|
||||||
|
all_hidden_states = all_hidden_states + (hidden_states,)
|
||||||
|
|
||||||
|
# Final layer norm.
|
||||||
|
if self.post_layer_norm:
|
||||||
|
hidden_states = self.final_layernorm(hidden_states)
|
||||||
|
|
||||||
|
return hidden_states, presents, all_hidden_states, all_self_attentions
|
||||||
|
|
||||||
|
|
||||||
def chatglm2_attention_forward(
|
def chatglm2_attention_forward(
|
||||||
self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
|
self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
|
||||||
):
|
):
|
||||||
|
|
|
||||||
|
|
@ -1,206 +0,0 @@
|
||||||
#
|
|
||||||
# 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-32k/blob/main/configuration_chatglm.py
|
|
||||||
#
|
|
||||||
|
|
||||||
import torch
|
|
||||||
from typing import Optional, Tuple, Union, List, Callable, Dict, Any
|
|
||||||
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.chatglm2 import core_attn_forward_8eb45c
|
|
||||||
|
|
||||||
|
|
||||||
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(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 chatglm2_32k_attention_forward(
|
|
||||||
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)
|
|
||||||
|
|
||||||
# apply relative positional encoding (rotary embedding)
|
|
||||||
if rotary_pos_emb is not None:
|
|
||||||
query_layer = apply_rotary_pos_emb(query_layer, rotary_pos_emb)
|
|
||||||
key_layer = apply_rotary_pos_emb(key_layer, rotary_pos_emb)
|
|
||||||
|
|
||||||
cur_length, batch_size = query_layer.shape[0], query_layer.shape[1]
|
|
||||||
|
|
||||||
if self.multi_query_attention:
|
|
||||||
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)
|
|
||||||
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
|
|
||||||
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
|
|
||||||
key_cache, value_cache = init_kv_cache(batch_size, self.num_attention_heads_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 use_cache:
|
|
||||||
key_layer = key_layer.permute(2, 0, 1, 3)
|
|
||||||
value_layer = value_layer.permute(2, 0, 1, 3)
|
|
||||||
if kv_cache is None:
|
|
||||||
kv_cache = torch.cat((key_layer.unsqueeze(0).unsqueeze(0),
|
|
||||||
value_layer.unsqueeze(0).unsqueeze(0)), dim=1)
|
|
||||||
else:
|
|
||||||
kv_cache = (key_layer, value_layer)
|
|
||||||
else:
|
|
||||||
kv_cache = None
|
|
||||||
|
|
||||||
# ==================================
|
|
||||||
# core attention computation
|
|
||||||
# ==================================
|
|
||||||
|
|
||||||
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, kv_cache
|
|
||||||
Loading…
Reference in a new issue