ayo/ipex-llm
1
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Chatglm support compresskv (#11690)

Browse source 
* chatglm4 support compresskv

* fix

* fix style

* support chatglm2

* fix quantkv conflict

* fix style
This commit is contained in:
Yina Chen 2024-08-01 13:20:20 +03:00 committed by GitHub
parent 762ad49362
commit 45c730ff39
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 117 additions and 42 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
python/llm/src/ipex_llm/transformers/models/chatglm2.py
View file

@ -25,6 +25,9 @@ from ipex_llm.utils.common.log4Error import invalidInputError
from ipex_llm.transformers.models.utils import restore_fp8_kv_cache, update_past_key_value from ipex_llm.transformers.models.utils import restore_fp8_kv_cache, update_past_key_value
from ipex_llm.transformers.models.utils import use_quantize_kv_cache, use_sdp, use_sdp_causal from ipex_llm.transformers.models.utils import use_quantize_kv_cache, use_sdp, use_sdp_causal
from ipex_llm.transformers.models.utils import should_use_fuse_rope, apply_rotary_pos_emb from ipex_llm.transformers.models.utils import should_use_fuse_rope, apply_rotary_pos_emb
from ipex_llm.transformers.models.utils import use_quantize_kv_cache, use_sdp, \
use_sdp_causal, should_use_compresskv, is_enough_kv_cache_room_4_36
from ipex_llm.transformers.kv import DynamicCompressCache
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
@ -83,6 +86,14 @@ def chatglm2_model_forward(
input_ids = torch.empty((batch_size, seq_length), input_ids = torch.empty((batch_size, seq_length),
dtype=inputs_embeds.dtype, device=inputs_embeds.device) dtype=inputs_embeds.dtype, device=inputs_embeds.device)
if use_cache:
use_compress_kv = should_use_compresskv(input_ids)
use_quantize_kv = use_quantize_kv_cache(self.encoder.layers[0].mlp.dense_h_to_4h,
input_ids)
if use_compress_kv and not use_quantize_kv and not isinstance(past_key_values,
DynamicCompressCache):
past_key_values = DynamicCompressCache.from_legacy_cache(past_key_values)
if full_attention_mask is None: if full_attention_mask is None:
if (attention_mask is not None and not attention_mask.all()) or ( if (attention_mask is not None and not attention_mask.all()) or (
past_key_values and seq_length != 1): past_key_values and seq_length != 1):
@ -157,7 +168,10 @@ def chatglm2_encoder_forward(
use_cache: Optional[bool] = True, use_cache: Optional[bool] = True,
output_hidden_states: Optional[bool] = False, output_hidden_states: Optional[bool] = False,
): ):
if not kv_caches: # [CompressKV]
use_compress_kv = isinstance(kv_caches, DynamicCompressCache)
if not kv_caches and not use_compress_kv:
kv_caches = [None for _ in range(self.num_layers)] kv_caches = [None for _ in range(self.num_layers)]
presents = () if use_cache else None presents = () if use_cache else None
if self.gradient_checkpointing and self.training: if self.gradient_checkpointing and self.training:
@ -184,11 +198,14 @@ def chatglm2_encoder_forward(
hidden_states, hidden_states,
attention_mask, attention_mask,
rotary_pos_emb, rotary_pos_emb,
kv_cache=kv_caches[index], kv_cache=kv_caches if use_compress_kv else kv_caches[index],
use_cache=use_cache use_cache=use_cache
) )
hidden_states, kv_cache = layer_ret hidden_states, kv_cache = layer_ret
if use_cache: if use_cache:
if use_compress_kv:
presents = kv_caches
else:
presents = presents + (kv_cache,) presents = presents + (kv_cache,)
if output_hidden_states: if output_hidden_states:
@ -207,8 +224,14 @@ def chatglm2_attention_forward(
# hidden_states: [seq_len, bsz, head_dim] # hidden_states: [seq_len, bsz, head_dim]
q_len, bsz, _ = hidden_states.size() q_len, bsz, _ = hidden_states.size()
# [CompressKV]
use_compresskv = isinstance(kv_cache, DynamicCompressCache)
# kv_cache: [seq_len, bsz, n_kv_head, head_dim] -> # kv_cache: [seq_len, bsz, n_kv_head, head_dim] ->
# past_key_value: [bsz, n_kv_head, seq_len, head_dim] # past_key_value: [bsz, n_kv_head, seq_len, head_dim]
if use_compresskv:
past_key_value = kv_cache
else:
past_key_value = None if kv_cache is None else (kv_cache[0].permute(1, 2, 0, 3), past_key_value = None if kv_cache is None else (kv_cache[0].permute(1, 2, 0, 3),
kv_cache[1].permute(1, 2, 0, 3)) kv_cache[1].permute(1, 2, 0, 3))
@ -227,6 +250,10 @@ def chatglm2_attention_forward(
kv_seq_len = key_states.shape[2] kv_seq_len = key_states.shape[2]
if past_key_value is not None: if past_key_value is not None:
if use_compresskv:
kv_seq_len += past_key_value.get_usable_length(kv_seq_len,
self.layer_number - 1)
else:
kv_seq_len += past_key_value[0].shape[2] kv_seq_len += past_key_value[0].shape[2]
# IPEX-LLM OPT: fuse rope # IPEX-LLM OPT: fuse rope
@ -249,6 +276,7 @@ def chatglm2_attention_forward(
# IPEX-LLM OPT: kv cache and quantize kv # IPEX-LLM OPT: kv cache and quantize kv
use_quantize_kv = use_quantize_kv_cache(self.query_key_value, query_states) use_quantize_kv = use_quantize_kv_cache(self.query_key_value, query_states)
if use_quantize_kv or (not use_compresskv):
key_states, value_states = update_past_key_value( key_states, value_states = update_past_key_value(
past_key_value, key_states, value_states, past_key_value, key_states, value_states,
kv_seq_len, use_quantize_kv, hidden_states.device kv_seq_len, use_quantize_kv, hidden_states.device
@ -256,6 +284,15 @@ def chatglm2_attention_forward(
# past_key_value: [bsz, n_kv_head, seq_len, head_dim] -> [seq_len, bsz, n_kv_head, head_dim] # past_key_value: [bsz, n_kv_head, seq_len, head_dim] -> [seq_len, bsz, n_kv_head, head_dim]
past_key_value = (key_states.permute(2, 0, 1, 3), past_key_value = (key_states.permute(2, 0, 1, 3),
value_states.permute(2, 0, 1, 3)) if use_cache else None value_states.permute(2, 0, 1, 3)) if use_cache else None
else:
from transformers.configuration_utils import PretrainedConfig
self.config = self.config if hasattr(self, "config") else PretrainedConfig()
enough_kv_room = is_enough_kv_cache_room_4_36(past_key_value, self.layer_number - 1)
key_states, value_states = past_key_value.update(
key_states, value_states, self.layer_number - 1,
query_states, attention_mask, n_head // n_kv_head,
self.config, enough_kv_room, 256
)
# IPEX-LLM OPT: sdp # IPEX-LLM OPT: sdp
attn_weights = None attn_weights = None

48
python/llm/src/ipex_llm/transformers/models/chatglm4.py
View file

@ -20,9 +20,11 @@
import torch import torch
from typing import Optional, Tuple, Union from typing import Optional, Tuple, Union
from ipex_llm.transformers.models.utils import restore_fp8_kv_cache, update_past_key_value from ipex_llm.transformers.models.utils import restore_fp8_kv_cache, update_past_key_value
from ipex_llm.transformers.models.utils import use_quantize_kv_cache, use_sdp, use_sdp_causal from ipex_llm.transformers.models.utils import use_quantize_kv_cache, use_sdp, \
use_sdp_causal, should_use_compresskv, is_enough_kv_cache_room_4_36
from ipex_llm.transformers.models.utils import should_use_fuse_rope, apply_rotary_pos_emb from ipex_llm.transformers.models.utils import should_use_fuse_rope, apply_rotary_pos_emb
from ipex_llm.transformers.models.chatglm2 import repeat_kv from ipex_llm.transformers.models.chatglm2 import repeat_kv
from ipex_llm.transformers.kv import DynamicCompressCache
from transformers.modeling_outputs import BaseModelOutputWithPast from transformers.modeling_outputs import BaseModelOutputWithPast
import math import math
@ -46,6 +48,15 @@ def chatglm4_model_forward(
use_cache = use_cache if use_cache is not None else self.config.use_cache 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 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if use_cache:
inputs = input_ids if input_ids is not None else inputs_embeds
use_compress_kv = should_use_compresskv(inputs)
use_quantize_kv = use_quantize_kv_cache(self.encoder.layers[0].mlp.dense_h_to_4h,
inputs)
if use_compress_kv and not use_quantize_kv and not isinstance(past_key_values,
DynamicCompressCache):
past_key_values = DynamicCompressCache.from_legacy_cache(past_key_values)
if inputs_embeds is None: if inputs_embeds is None:
batch_size, seq_length = input_ids.shape batch_size, seq_length = input_ids.shape
inputs_embeds = self.embedding(input_ids) inputs_embeds = self.embedding(input_ids)
@ -134,7 +145,13 @@ def chatglm4_attention_forward(
# hidden_states: [b, sq, h] # hidden_states: [b, sq, h]
bsz, q_len, _ = hidden_states.size() bsz, q_len, _ = hidden_states.size()
# [CompressKV]
use_compresskv = isinstance(kv_cache, DynamicCompressCache)
# past_key_value: [bsz, n_kv_head, seq_len, head_dim] # past_key_value: [bsz, n_kv_head, seq_len, head_dim]
if use_compresskv:
past_key_value = kv_cache
else:
past_key_value = None if kv_cache is None else (kv_cache[0], past_key_value = None if kv_cache is None else (kv_cache[0],
kv_cache[1]) kv_cache[1])
@ -153,6 +170,10 @@ def chatglm4_attention_forward(
kv_seq_len = key_states.shape[2] kv_seq_len = key_states.shape[2]
if past_key_value is not None: if past_key_value is not None:
if use_compresskv:
kv_seq_len += past_key_value.get_usable_length(kv_seq_len,
self.layer_number - 1)
else:
kv_seq_len += past_key_value[0].shape[2] kv_seq_len += past_key_value[0].shape[2]
# IPEX-LLM OPT: fuse rope # IPEX-LLM OPT: fuse rope
@ -175,11 +196,12 @@ def chatglm4_attention_forward(
# IPEX-LLM OPT: kv cache and quantize kv # IPEX-LLM OPT: kv cache and quantize kv
use_quantize_kv = use_quantize_kv_cache(self.query_key_value, query_states) use_quantize_kv = use_quantize_kv_cache(self.query_key_value, query_states)
if use_quantize_kv or (not use_compresskv):
key_states, value_states = update_past_key_value( key_states, value_states = update_past_key_value(
past_key_value, key_states, value_states, past_key_value, key_states, value_states,
kv_seq_len, use_quantize_kv, hidden_states.device kv_seq_len, use_quantize_kv, hidden_states.device
) )
if use_cache: if use_cache:
if past_key_value is None: if past_key_value is None:
past_key_value = torch.cat((key_states.unsqueeze(0).unsqueeze(0), past_key_value = torch.cat((key_states.unsqueeze(0).unsqueeze(0),
@ -188,6 +210,15 @@ def chatglm4_attention_forward(
past_key_value = (key_states, value_states) past_key_value = (key_states, value_states)
else: else:
past_key_value = None past_key_value = None
else:
from transformers.configuration_utils import PretrainedConfig
self.config = self.config if hasattr(self, "config") else PretrainedConfig()
enough_kv_room = is_enough_kv_cache_room_4_36(past_key_value, self.layer_number - 1)
key_states, value_states = past_key_value.update(
key_states, value_states, self.layer_number - 1,
query_states, attention_mask, n_head // n_kv_head,
self.config, enough_kv_room, 256
)
# IPEX-LLM OPT: sdp # IPEX-LLM OPT: sdp
attn_weights = None attn_weights = None
@ -244,7 +275,10 @@ def chatglm4_encoder_forward(
use_cache: Optional[bool] = True, use_cache: Optional[bool] = True,
output_hidden_states: Optional[bool] = False, output_hidden_states: Optional[bool] = False,
): ):
if not kv_caches: # [CompressKV]
use_compress_kv = isinstance(kv_caches, DynamicCompressCache)
if not kv_caches and not use_compress_kv:
kv_caches = [None for _ in range(self.num_layers)] kv_caches = [None for _ in range(self.num_layers)]
presents = () if use_cache else None presents = () if use_cache else None
if self.gradient_checkpointing and self.training: if self.gradient_checkpointing and self.training:
@ -274,11 +308,14 @@ def chatglm4_encoder_forward(
hidden_states, hidden_states,
attention_mask, attention_mask,
rotary_pos_emb, rotary_pos_emb,
kv_cache=kv_caches[index], kv_cache=kv_caches if use_compress_kv else kv_caches[index],
use_cache=use_cache use_cache=use_cache
) )
hidden_states, kv_cache = layer_ret hidden_states, kv_cache = layer_ret
if use_cache: if use_cache:
if use_compress_kv:
presents = kv_caches
else:
# token by token decoding, use tuple format # token by token decoding, use tuple format
if kv_caches[0] is not None: if kv_caches[0] is not None:
presents = presents + (kv_cache,) presents = presents + (kv_cache,)
@ -290,7 +327,8 @@ def chatglm4_encoder_forward(
# bigdl-llm change starts # bigdl-llm change starts
# to fix first token's kv cache error of tensor format in pipeline parallel # to fix first token's kv cache error of tensor format in pipeline parallel
if isinstance(kv_cache, tuple): if isinstance(kv_cache, tuple):
kv_cache = torch.tensor(kv_cache, kv_cache = torch.tensor(
kv_cache,
dtype=hidden_states.dtype).to(hidden_states.device) dtype=hidden_states.dtype).to(hidden_states.device)
# bigdl-llm change ends # bigdl-llm change ends
presents = torch.cat((presents, kv_cache.to(presents.device)), dim=0) presents = torch.cat((presents, kv_cache.to(presents.device)), dim=0)