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 mllama, gpt2 and internvl #12602

Merged
merged 1 commit into from
Dec 24, 2024
Merged

refactor mllama, gpt2 and internvl #12602

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
8 changes: 5 additions & 3 deletions python/llm/src/ipex_llm/transformers/models/gpt2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
#

import torch
from ipex_llm.transformers.models.common import scaled_dot_product_attention
from ipex_llm.transformers.models.utils import use_sdp_non_causal


Expand Down Expand Up @@ -44,10 +45,11 @@ def gpt2_attention_attn(
else:
attention_mask = attention_mask.expand(-1, -1, seq_len, seq_len)

import xe_addons
attn_weights = None
attn_output = xe_addons.sdp_non_causal(query, key.contiguous(),
value.contiguous(), attention_mask)
attn_output = scaled_dot_product_attention(
query, key.contiguous(), value.contiguous(),
attention_mask, False
)
return attn_output, attn_weights
# ipex-llm changes end

Expand Down
7 changes: 5 additions & 2 deletions python/llm/src/ipex_llm/transformers/models/internvl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@

import torch
from ipex_llm.utils.common.log4Error import invalidInputError
from ipex_llm.transformers.models.common import scaled_dot_product_attention
from ipex_llm.transformers.models.utils import use_sdp_non_causal


Expand Down Expand Up @@ -177,8 +178,10 @@ def intern_attention_forward(self, x: torch.Tensor) -> torch.Tensor:
k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)

if use_sdp_non_causal(self.head_dim, q.device, q.dtype):
import xe_addons
x = xe_addons.sdp_non_causal(q, k.contiguous(), v.contiguous(), None)
x = scaled_dot_product_attention(
q, k.contiguous(), v.contiguous(),
None, False, self.scale
)
else:
attn = ((q * self.scale) @ k.transpose(-2, -1))
attn = attn.softmax(dim=-1)
Expand Down
73 changes: 11 additions & 62 deletions python/llm/src/ipex_llm/transformers/models/mllama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,19 +32,17 @@
# limitations under the License.


import math
import torch

from typing import Optional, Tuple, Union
from transformers.cache_utils import Cache
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.models.mllama.modeling_mllama import MllamaVisionAttention
from transformers.models.mllama.modeling_mllama import MllamaTextSelfAttention
from transformers.models.mllama.modeling_mllama import repeat_kv
from ipex_llm.transformers.models.utils import use_sdp, use_sdp_causal, use_sdp_non_causal
from ipex_llm.transformers.models.utils import use_quantize_kv_cache, restore_fp8_kv_cache
from ipex_llm.transformers.models.utils import use_quantize_kv_cache
from ipex_llm.transformers.models.utils import should_use_fuse_rope
from ipex_llm.transformers.models.common import merge_qkv_base, attention_softmax
from ipex_llm.transformers.models.common import scaled_dot_product_attention
from ipex_llm.transformers.kv import DynamicNormalCache, DynamicFp8Cache
from ipex_llm.transformers.utils import invalidInputError

Expand All @@ -67,27 +65,11 @@ def mllama_vision_attention_forward(
qkv = qkv.transpose(1, 2)
query, key, value = qkv.chunk(3, dim=1)

if attention_mask is not None: # no matter the length, we just slice it
causal_mask = attention_mask[:, :, :, : key.shape[-2]]
else:
causal_mask = None

if use_sdp_non_causal(self.head_dim, query.device, query.dtype):
import xe_addons
attn_output = xe_addons.sdp_non_causal(query, key.contiguous(),
value.contiguous(), causal_mask)
attn_weights = None
else:
attn_weights = torch.matmul(query, key.transpose(2, 3)) / math.sqrt(self.head_dim)

if attention_mask is not None:
attn_weights = attn_weights + causal_mask

# upcast attention to fp32
from ipex_llm.transformers.models.common import attention_softmax
attn_weights = attention_softmax(attn_weights)

attn_output = torch.matmul(attn_weights, value)
attn_weights = None
attn_output = scaled_dot_product_attention(
query, key.contiguous(), value.contiguous(),
attention_softmax, False
)

attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, -1)
Expand Down Expand Up @@ -278,44 +260,11 @@ def mllama_cross_attention_forward(
past_key_value.value_cache[self.layer_idx],
)

kv_seq_len = key_states.size(2)
if attention_mask is not None: # no matter the length, we just slice it
causal_mask = attention_mask[:, :, :, :kv_seq_len]
else:
causal_mask = None

attn_weights = None
if use_sdp(q_len, kv_seq_len, self.head_dim, query_states):
import xe_addons
if isinstance(past_key_value, DynamicFp8Cache):
attn_output = xe_addons.sdp_fp8(query_states, key_states, value_states, causal_mask)
else:
attn_output = xe_addons.sdp(query_states, key_states, value_states, causal_mask)
elif use_sdp_causal(q_len, kv_seq_len, self.head_dim, query_states, self.training):
import xe_addons
if isinstance(past_key_value, DynamicFp8Cache):
attn_output = xe_addons.sdp_fp8_causal(query_states, key_states,
value_states, causal_mask)
else:
attn_output = xe_addons.sdp_causal(query_states, key_states,
value_states, causal_mask)
else:
if isinstance(past_key_value, DynamicFp8Cache):
key_states, value_states = restore_fp8_kv_cache(key_states, value_states,
query_states.dtype)
# repeat k/v heads if n_kv_heads < n_heads
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)

attn_weights = torch.matmul(query_states,
key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

if causal_mask is not None:
attn_weights = attn_weights + causal_mask

# upcast attention to fp32
attn_weights = attention_softmax(attn_weights)
attn_output = torch.matmul(attn_weights, value_states)
attn_output = scaled_dot_product_attention(
query_states, key_states, value_states,
attention_mask, q_len == key_states.size(2)
)

attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, -1)
Expand Down
Loading