Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor codegeex to remove ipex kernel usage #12664

Merged
merged 2 commits into from
Jan 7, 2025
Merged

refactor codegeex to remove ipex kernel usage #12664

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
feb550f
refactor codegeex to remove ipex kernel usage
MeouSker77 Jan 7, 2025
35fc475
small update
MeouSker77 Jan 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 6 additions & 23 deletions python/llm/src/ipex_llm/transformers/convert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1052,7 +1052,8 @@ def _optimize_pre(model, qtype=None):
_optimize_pre(model.llm, qtype=qtype)
model.llm.config.model_type = "megrezo"
elif model.config.model_type == "chatglm":
if hasattr(model.config, 'padded_vocab_size') and model.config.padded_vocab_size == 65024:
if hasattr(model.config, 'padded_vocab_size') and \
model.config.padded_vocab_size in [65024, 64896]:
# chatglm2 and chatglm3
from ipex_llm.transformers.models.chatglm2 import split_mlp
model.apply(split_mlp)
Expand Down Expand Up @@ -1337,7 +1338,7 @@ def _optimize_post(model):
and model.config.architectures[0] in ["ChatGLMModel", "ChatGLMForConditionalGeneration"]
):
if hasattr(model.config, 'padded_vocab_size') and \
model.config.padded_vocab_size == 65024:
model.config.padded_vocab_size in [65024, 64896]:
# chatglm2-6b, chatglm2-6b-32k, chatglm3-6b, chatglm3-6b-32k, chatglm3-6b-128k
modeling_module_name = model.__class__.__module__
module = importlib.import_module(modeling_module_name)
Expand All @@ -1359,27 +1360,9 @@ def _optimize_post(model):
module.RMSNorm,
chatglm_rms_norm_forward)
convert_forward(model, module.MLP, mlp_forward)
elif hasattr(model.config, 'padded_vocab_size') and \
model.config.padded_vocab_size == 64896:
# codegeex-nano
modeling_module_name = model.__class__.__module__
module = importlib.import_module(modeling_module_name)
from ipex_llm.transformers.models.chatglm2 import codegeex_attention_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 codegeex_model_forward
convert_forward(model,
module.SelfAttention,
codegeex_attention_forward)
convert_forward(model,
module.GLMTransformer,
chatglm2_encoder_forward)
convert_forward(model,
module.ChatGLMModel,
codegeex_model_forward)
convert_forward(model,
module.RMSNorm,
chatglm_rms_norm_forward)
# for codegeex-nano
if hasattr(model.config, "rope_ratio"):
model.transformer.rotary_pos_emb.rope_ratio = model.config.rope_ratio
elif hasattr(model.config, 'vocab_size') and model.config.vocab_size == 130528:
# chatglm-6b
modeling_module_name = model.__class__.__module__
Expand Down
193 changes: 1 addition & 192 deletions python/llm/src/ipex_llm/transformers/models/chatglm2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,7 @@ def chatglm2_attention_forward(
# IPEX-LLM OPT: fuse rope
inv_freq, position_ids = rotary_pos_emb
rot_dim = inv_freq.size(-1) * 2
if should_use_fuse_rope(hidden_states, rotary_pos_emb[1], self.training):
if should_use_fuse_rope(hidden_states, position_ids, self.training):
import xe_addons
xe_addons.rotary_two_inplaced(inv_freq, position_ids,
query_states[..., :rot_dim], key_states[..., :rot_dim])
Expand Down Expand Up @@ -321,197 +321,6 @@ def chatglm2_attention_forward(
return output, past_key_value


@torch.jit.script
def apply_rotary_pos_emb_original(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 codegeex_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

if inputs_embeds is None:
batch_size, seq_length = input_ids.shape
inputs_embeds = self.embedding(input_ids)
else:
inputs_embeds = inputs_embeds.transpose(0, 1).contiguous()
seq_length, batch_size, _ = inputs_embeds.shape
input_ids = torch.empty((batch_size, seq_length),
dtype=inputs_embeds.dtype, device=inputs_embeds.device)

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)

# ipex-llm changes begin
# 1. replace `rotary_pos_emb` with `inv_freq` and `position_ids`
# 2. generate `causal_mask` and replace `full_attention_mask` with it
if position_ids is None:
if past_key_values is None:
position_ids = torch.arange(seq_length, dtype=torch.int64, device=inputs_embeds.device)
else:
if isinstance(past_key_values, DynamicCompressCache):
kv_length = past_key_values.get_seq_length()
else:
kv_length = past_key_values[0][0].size(0)
position_ids = torch.arange(kv_length, kv_length + seq_length,
dtype=torch.int64, device=inputs_embeds.device)
position_ids = position_ids.repeat(batch_size, 1)
use_fuse_rope = input_ids.device.type == "xpu" 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()

# `full_attention_mask` is not None only when
# `past_key_values` is not None and `seq_length` > 1
if full_attention_mask is not None:
causal_mask = torch.zeros([batch_size, 1, seq_length, full_attention_mask.size(-1)],
dtype=inputs_embeds.dtype, device=inputs_embeds.device)
mask_value = torch.finfo(inputs_embeds.dtype).min
causal_mask.masked_fill_(full_attention_mask, mask_value)
elif self.training or (inputs_embeds.device.type != "xpu" and past_key_values is None):
full_attention_mask = self.get_masks(input_ids,
past_key_values,
padding_mask=attention_mask)
causal_mask = torch.zeros([batch_size, 1, seq_length, full_attention_mask.size(-1)],
dtype=inputs_embeds.dtype, device=inputs_embeds.device)
mask_value = torch.finfo(inputs_embeds.dtype).min
causal_mask.masked_fill_(full_attention_mask, mask_value)
else:
causal_mask = None

# Run encoder.
hidden_states, presents, all_hidden_states, all_self_attentions = self.encoder(
inputs_embeds, causal_mask,
rotary_pos_emb=rotary_pos_emb,
kv_caches=past_key_values, use_cache=use_cache, output_hidden_states=output_hidden_states
)
# ipex-llm changes end

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 codegeex_attention_forward(
self, hidden_states, attention_mask, rotary_pos_emb, kv_cache=None, use_cache=True
):
q_len, bsz, _ = hidden_states.size()
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

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))
qkv = self.query_key_value(hidden_states)
qkv = qkv.view(q_len, bsz, n_head + 2 * n_kv_head, head_dim)
# [seq_len, bsz, n_head, head_dim] -> [bsz, n_head, seq_len, head_dim]
qkv = qkv.permute(1, 2, 0, 3)
query_layer, key_layer, value_layer = qkv.split([n_head,
n_kv_head,
n_kv_head], dim=1)
kv_seq_len = key_layer.shape[2]
if past_key_value is not None:
kv_seq_len += past_key_value[0].shape[2]

# apply relative positional encoding (rotary embedding)
if len(rotary_pos_emb) == 2 and isinstance(rotary_pos_emb, tuple):
cos, sin = rotary_pos_emb
rot_dim = cos.shape[-1]
query_layer = query_layer.transpose(1, 2)
key_layer = key_layer.transpose(1, 2)
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(1, 2)
key_layer = key_layer.transpose(1, 2)
else:
query_layer = apply_rotary_pos_emb_original(query_layer, rotary_pos_emb)
key_layer = apply_rotary_pos_emb_original(key_layer, rotary_pos_emb)

key_layer, value_layer = update_past_key_value(
past_key_value, key_layer, value_layer,
kv_seq_len, False, hidden_states.device
)
# past_key_value: [bsz, n_kv_head, seq_len, head_dim] -> [seq_len, bsz, n_kv_head, head_dim]
past_key_value = (key_layer.permute(2, 0, 1, 3),
value_layer.permute(2, 0, 1, 3)) if use_cache else None

# =================
# Output. [sq, b, h]
# =================
context_layer = scaled_dot_product_attention(
query_layer, key_layer, value_layer,
attention_mask, q_len == kv_seq_len
)

context_layer = context_layer.permute(2, 0, 1, 3).contiguous().view(q_len,
bsz,
n_head * head_dim)
output = self.dense(context_layer)

return output, past_key_value

import torch.nn.functional as F


Expand Down
Loading