glm 4v 1st sdp for vision (#12904)
* glm4v 1st sdp * update glm4v example * meet code review * fix style
This commit is contained in:
parent
5c100ac105
commit
e946127613
2 changed files with 14 additions and 22 deletions
|
|
@ -61,7 +61,7 @@ if __name__ == '__main__':
|
||||||
trust_remote_code=True,
|
trust_remote_code=True,
|
||||||
use_cache=True,
|
use_cache=True,
|
||||||
model_hub=model_hub)
|
model_hub=model_hub)
|
||||||
model = model.half().to('xpu')
|
model = model.to('xpu')
|
||||||
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -19,7 +19,7 @@
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from typing import Optional, Tuple, Union
|
from typing import Optional, Tuple, Union
|
||||||
from ipex_llm.transformers.models.common import merge_qkv_base
|
from ipex_llm.transformers.models.common import merge_qkv_base, padding_qkv_hd
|
||||||
from ipex_llm.transformers.models.common import scaled_dot_product_attention
|
from ipex_llm.transformers.models.common import scaled_dot_product_attention
|
||||||
from ipex_llm.transformers.models.utils import update_past_key_value
|
from ipex_llm.transformers.models.utils import update_past_key_value
|
||||||
from ipex_llm.transformers.models.utils import use_quantize_kv_cache, use_sdp
|
from ipex_llm.transformers.models.utils import use_quantize_kv_cache, use_sdp
|
||||||
|
|
@ -265,26 +265,18 @@ def visual_attention_forward(self, x: "tensor(B, L, D)") -> "tensor(B, L, D)":
|
||||||
q, k, v = qkv[0], qkv[1], qkv[2]
|
q, k, v = qkv[0], qkv[1], qkv[2]
|
||||||
|
|
||||||
bsz, q_len, kv_seq_len, head_dim = q.shape
|
bsz, q_len, kv_seq_len, head_dim = q.shape
|
||||||
if use_sdp(q_len, kv_seq_len, head_dim, q):
|
q, k, v = padding_qkv_hd(
|
||||||
import xe_addons
|
q, k, v,
|
||||||
out = xe_addons.sdp(q, k, v, None)
|
head_dim, 128
|
||||||
elif q.device.type == "cpu":
|
)
|
||||||
out = torch.nn.functional.scaled_dot_product_attention(q, k, v,
|
|
||||||
attn_mask=None,
|
attn_weights = None
|
||||||
dropout_p=0.,
|
attn_output = scaled_dot_product_attention(
|
||||||
is_causal=False)
|
q, k.contiguous(), v.contiguous(),
|
||||||
else:
|
None, False, 1 / math.sqrt(head_dim)
|
||||||
attn_weights = torch.matmul(q / math.sqrt(head_dim),
|
)
|
||||||
k.transpose(2, 3)).to(v.dtype)
|
|
||||||
if kv_seq_len >= 2048 or bsz >= 64:
|
out = attn_output[:, :, :, :head_dim]
|
||||||
# for memory considerations, do not upcast attention to fp32
|
|
||||||
# for long sequences or large batches
|
|
||||||
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
|
|
||||||
else:
|
|
||||||
# upcast attention to fp32
|
|
||||||
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1,
|
|
||||||
dtype=torch.float32).to(v.dtype)
|
|
||||||
out = torch.matmul(attn_weights, v)
|
|
||||||
output = self.dense(out.transpose(1, 2).reshape(B, L, -1))
|
output = self.dense(out.transpose(1, 2).reshape(B, L, -1))
|
||||||
output = self.output_dropout(output)
|
output = self.output_dropout(output)
|
||||||
return output
|
return output
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue