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

GPT-J rope optimization on xpu (#10182)

Browse source 
* optimize

* update

* fix style & move use_fuse_rope

* add ipex version check

* fix style

* update

* fix style

* meet comments

* address comments

* fix style
This commit is contained in:
Yina Chen 2024-02-22 16:25:12 +08:00 committed by GitHub
parent f445217d02
commit ce5840a8b7
4 changed files with 301 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
python/llm/src/bigdl/llm/transformers/convert.py
View file

@ -776,10 +776,17 @@ def _optimize_post(model, lightweight_bmm=False):
# dolly-v1-6b
modeling_module_name = model.__class__.__module__
module = importlib.import_module(modeling_module_name)
from bigdl.llm.transformers.models.gptj import gptj_attention_forward
from bigdl.llm.transformers.models.gptj import gptj_attention_forward, gptj_model_forward,\
gptj_block_forward
convert_forward(model,
module.GPTJAttention,
gptj_attention_forward)
convert_forward(model,
module.GPTJModel,
gptj_model_forward)
convert_forward(model,
module.GPTJBlock,
gptj_block_forward)
elif "bloom" in model.config.model_type:
modeling_module_name = model.__class__.__module__
module = importlib.import_module(modeling_module_name)

7
python/llm/src/bigdl/llm/transformers/model.py
View file

@ -90,6 +90,12 @@ def save_low_bit(self, *args, **kwargs):
self.to(origin_device)
def _load_pre():
from transformers import GPTJModel
from bigdl.llm.transformers.models.gptj import gptj_model_new_init
GPTJModel.__init__ = gptj_model_new_init
class _BaseAutoModelClass:
HF_MODEL = None
@ -399,6 +405,7 @@ class _BaseAutoModelClass:
offload_dir=None
)
else:
_load_pre()
try:
model = cls.HF_Model.from_pretrained(*args, **kwargs)
except NotImplementedError:

292
python/llm/src/bigdl/llm/transformers/models/gptj.py
View file

@ -20,8 +20,11 @@
import torch
from typing import Optional, Tuple, Union
from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, \
apply_rotary_pos_emb, append_kv_cache
apply_rotary_pos_emb, append_kv_cache, apply_ipex_rotate_every_two
from transformers.utils.import_utils import is_torch_fx_proxy
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.gptj.modeling_gptj import GPTJModel
from bigdl.llm.utils.common import invalidInputError
KV_CACHE_ALLOC_BLOCK_LENGTH = 256
@ -87,6 +90,7 @@ def gptj_attention_forward(
position_ids: Optional[torch.LongTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
use_cache: Optional[bool] = False,
rotary_emb: Optional[Tuple]=None,
output_attentions: Optional[bool] = False,
) -> Union[
Tuple[torch.Tensor, Tuple[torch.Tensor]],
@ -100,30 +104,26 @@ def gptj_attention_forward(
key = self._split_heads(key, self.num_attention_heads, self.head_dim, True)
value = self._split_heads(value, self.num_attention_heads, self.head_dim, False)
if is_torch_fx_proxy(position_ids) or torch.jit.is_tracing():
# The logic to conditionally copy to GPU could not be traced, so we do this
# every time in the torch.fx case
embed_positions = get_embed_positions(self.embed_positions, position_ids)
else:
embed_positions = self._get_embed_positions(position_ids)
repeated_position_ids = position_ids.unsqueeze(-1).repeat(1, 1, embed_positions.shape[-1])
sincos = torch.gather(embed_positions, 1, repeated_position_ids)
sin, cos = torch.split(sincos, sincos.shape[-1] // 2, dim=-1)
sin, cos = rotary_emb
use_fuse_rope = hidden_states.device.type == "xpu" and not self.training
if self.rotary_dim is not None:
k_rot = key[:, :, :, : self.rotary_dim]
k_pass = key[:, :, :, self.rotary_dim:]
q_rot = query[:, :, :, : self.rotary_dim]
q_pass = query[:, :, :, self.rotary_dim:]
q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin, position_ids, "gptj")
key = torch.cat([k_rot, k_pass], dim=-1)
query = torch.cat([q_rot, q_pass], dim=-1)
if use_fuse_rope:
apply_ipex_rotate_every_two(q_rot, k_rot, cos, sin)
else:
k_pass = key[:, :, :, self.rotary_dim:]
q_pass = query[:, :, :, self.rotary_dim:]
q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin, position_ids, "gptj")
key = torch.cat([k_rot, k_pass], dim=-1)
query = torch.cat([q_rot, q_pass], dim=-1)
else:
query, key = apply_rotary_pos_emb(query, key, cos, sin, position_ids, "gptj")
if use_fuse_rope:
apply_ipex_rotate_every_two(query, key, cos, sin)
else:
query, key = apply_rotary_pos_emb(query, key, cos, sin, position_ids, "gptj")
batch_size, q_len, _ = hidden_states.size()
@ -184,3 +184,257 @@ def gptj_attention_forward(
outputs += (attn_weights,)
return outputs # a, present, (attentions)
def gptj_block_forward(
self,
hidden_states: Optional[torch.FloatTensor],
layer_past: Optional[Tuple[torch.Tensor]] = None,
attention_mask: Optional[torch.FloatTensor] = None,
position_ids: Optional[torch.LongTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
use_cache: Optional[bool] = False,
rotary_emb: Optional[Tuple]=None,
output_attentions: Optional[bool] = False,
) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:
residual = hidden_states
hidden_states = self.ln_1(hidden_states)
attn_outputs = self.attn(
hidden_states=hidden_states,
layer_past=layer_past,
attention_mask=attention_mask,
position_ids=position_ids,
head_mask=head_mask,
use_cache=use_cache,
rotary_emb=rotary_emb,
output_attentions=output_attentions,
)
attn_output = attn_outputs[0] # output_attn: a, present, (attentions)
outputs = attn_outputs[1:]
feed_forward_hidden_states = self.mlp(hidden_states)
hidden_states = attn_output + feed_forward_hidden_states + residual
if use_cache:
outputs = (hidden_states,) + outputs
else:
outputs = (hidden_states,) + outputs[1:]
return outputs # hidden_states, present, (attentions)
def create_sinusoidal_positions(num_pos: int, dim: int) -> torch.Tensor:
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2) / dim))
sinusoid_inp = torch.einsum("i , j -> i j",
torch.arange(num_pos, dtype=torch.float), inv_freq).float()
return torch.cat((torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)), dim=1)
old_init = GPTJModel.__init__
def gptj_model_new_init(self, config):
old_init(self, config)
embed_dim = config.hidden_size
rotary_dim = config.rotary_dim
pos_embd_dim = rotary_dim or embed_dim
max_positions = config.max_position_embeddings
self.embed_positions = create_sinusoidal_positions(max_positions, pos_embd_dim)
def get_new_embed_positions(position_ids, prev_embed_positions):
embed_positions = prev_embed_positions
if embed_positions.device != position_ids.device:
embed_positions = embed_positions.to(position_ids.device)
prev_embed_positions = embed_positions
return embed_positions.repeat(position_ids.shape[0], 1, 1), prev_embed_positions
def gptj_model_forward(
self,
input_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
attention_mask: Optional[torch.FloatTensor] = None,
token_type_ids: Optional[torch.LongTensor] = None,
position_ids: Optional[torch.LongTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
) -> Union[Tuple, BaseModelOutputWithPast]:
output_attentions = output_attentions if output_attentions is not None \
else self.config.output_attentions
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 input_ids is not None and inputs_embeds is not None:
invalidInputError(False,
"You cannot specify both input_ids and inputs_embeds at the same time")
elif input_ids is not None:
self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
input_shape = input_ids.size()
input_ids = input_ids.view(-1, input_shape[-1])
batch_size = input_ids.shape[0]
elif inputs_embeds is not None:
input_shape = inputs_embeds.size()[:-1]
batch_size = inputs_embeds.shape[0]
else:
invalidInputError(False, "You have to specify either input_ids or inputs_embeds")
device = input_ids.device if input_ids is not None else inputs_embeds.device
if token_type_ids is not None:
token_type_ids = token_type_ids.view(-1, input_shape[-1])
if past_key_values is None:
past_length = 0
past_key_values = tuple([None] * len(self.h))
else:
past_length = past_key_values[0][0].size(-2)
if position_ids is None:
position_ids = torch.arange(past_length, input_shape[-1] + past_length,
dtype=torch.long, device=device)
position_ids = position_ids.unsqueeze(0)
# Attention mask.
if attention_mask is not None:
if batch_size <= 0:
invalidInputError(False, "batch_size has to be defined and > 0")
attention_mask = attention_mask.view(batch_size, -1)
# We create a 3D attention mask from a 2D tensor mask.
# Sizes are [batch_size, 1, 1, to_seq_length]
# So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
# this attention mask is more simple than the triangular masking of causal attention
# used in OpenAI GPT, we just need to prepare the broadcast dimension here.
attention_mask = attention_mask[:, None, None, :]
# Since attention_mask is 1.0 for positions we want to attend and 0.0 for
# masked positions, this operation will create a tensor which is 0.0 for
# positions we want to attend and the dtype's smallest value for masked positions.
# Since we are adding it to the raw scores before the softmax, this is
# effectively the same as removing these entirely.
attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
# Prepare head mask if needed
# 1.0 in head_mask indicate we keep the head
# attention_probs has shape bsz x num_attention_heads x N x N
# head_mask has shape n_layer x batch x num_attention_heads x N x N
head_mask = self.get_head_mask(head_mask, self.config.n_layer)
if inputs_embeds is None:
inputs_embeds = self.wte(input_ids)
hidden_states = inputs_embeds
if token_type_ids is not None:
token_type_embeds = self.wte(token_type_ids)
hidden_states = hidden_states + token_type_embeds
hidden_states = self.drop(hidden_states)
output_shape = (-1,) + input_shape[1:] + (hidden_states.size(-1),)
if self.gradient_checkpointing and self.training:
if use_cache:
logger.warning_once(
"`use_cache=True` is incompatible with gradient checkpointing."
"Setting `use_cache=False`..."
)
use_cache = False
presents = () if use_cache else None
all_self_attentions = () if output_attentions else None
all_hidden_states = () if output_hidden_states else None
# Repeat cos sin here, call only once for each token.
# If put this to attension forward, it will generate too many times.
if is_torch_fx_proxy(position_ids) or torch.jit.is_tracing():
# The logic to conditionally copy to GPU could not be traced, so we do this
# every time in the torch.fx case
embed_positions = get_embed_positions(self.embed_positions, position_ids)
else:
embed_positions, self.embed_positions = get_new_embed_positions(position_ids,
self.embed_positions)
repeated_position_ids = position_ids.unsqueeze(-1).repeat(1, 1, embed_positions.shape[-1])
sincos = torch.gather(embed_positions, 1, repeated_position_ids)
sin, cos = torch.split(sincos, sincos.shape[-1] // 2, dim=-1)
sin = torch.repeat_interleave(sin[:, :, None, :], 2, 3)
cos = torch.repeat_interleave(cos[:, :, None, :], 2, 3)
for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
# Model parallel
if self.model_parallel:
torch.cuda.set_device(hidden_states.device)
# Ensure layer_past is on same device as hidden_states (might not be correct)
if layer_past is not None:
layer_past = tuple(past_state.to(hidden_states.device) for past_state in layer_past)
# Ensure that attention_mask is always on the same device as hidden_states
if attention_mask is not None:
attention_mask = attention_mask.to(hidden_states.device)
if isinstance(head_mask, torch.Tensor):
head_mask = head_mask.to(hidden_states.device)
if output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
if self.gradient_checkpointing and self.training:
outputs = self._gradient_checkpointing_func(
block.__call__,
hidden_states,
None,
attention_mask,
position_ids,
head_mask[i],
use_cache,
output_attentions,
)
else:
outputs = block(
hidden_states=hidden_states,
layer_past=layer_past,
attention_mask=attention_mask,
position_ids=position_ids,
head_mask=head_mask[i],
use_cache=use_cache,
rotary_emb=(sin, cos),
output_attentions=output_attentions,
)
hidden_states = outputs[0]
if use_cache is True:
presents = presents + (outputs[1],)
if output_attentions:
all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
# Model Parallel: If it's the last layer for that device, put things on the next device
if self.model_parallel:
for k, v in self.device_map.items():
if i == v[-1] and "cuda:" + str(k) != self.last_device:
hidden_states = hidden_states.to("cuda:" + str(k + 1))
hidden_states = self.ln_f(hidden_states)
hidden_states = hidden_states.view(output_shape)
# Add last hidden state
if output_hidden_states:
all_hidden_states = all_hidden_states + (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,
)

15
python/llm/src/bigdl/llm/transformers/models/utils.py
View file

@ -153,8 +153,6 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, model_family):
k_embed = (k * cos) + (rotate_half(k) * sin)
return q_embed, k_embed
elif model_family == "gptj":
cos = torch.repeat_interleave(cos[:, :, None, :], 2, 3)
sin = torch.repeat_interleave(sin[:, :, None, :], 2, 3)
q_embed = (q * cos) + (rotate_every_two(q) * sin)
k_embed = (k * cos) + (rotate_every_two(k) * sin)
return q_embed, k_embed
@ -163,6 +161,19 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, model_family):
f"{model_family} is not supported.")
def apply_ipex_rotate_every_two(q, k, cos, sin):
# ipex's apply_rotary_embedding_two_qk can change the origin storage,
# so q/k will get the result directly.
from bigdl.llm.transformers.utils import get_ipex_version
if get_ipex_version() >= "2.1.10+xpu":
torch.ops.torch_ipex.apply_rotary_embedding_two_qk(
q, k, sin, cos, q, k
)
else:
torch.ops.torch_ipex.apply_rotary_embedding(q, sin, cos, q)
torch.ops.torch_ipex.apply_rotary_embedding(k, sin, cos, k)
def apply_rotary_pos_emb_no_cache_xpu(q, k, position_ids, model_family):
if q.device.type != "xpu":
invalidInputError(False,