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

Baichuan 7b fp16 sdp and qwen2 pvc sdp (#10435)

Browse source 
* add baichuan sdp

* update

* baichuan2

* fix

* fix style

* revert 13b

* revert
This commit is contained in:
Xin Qiu 2024-03-18 10:15:34 +08:00 committed by GitHub
parent 5ab52ef5b5
commit 399843faf0
3 changed files with 77 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
python/llm/src/bigdl/llm/transformers/models/baichuan.py
View file

@ -24,8 +24,10 @@ from typing import List, Optional, Tuple, Union
import torch import torch
import torch.utils.checkpoint import torch.utils.checkpoint
from torch import nn from torch import nn
import torch.nn.functional as F
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
from bigdl.llm.utils.common import invalidInputError from bigdl.llm.utils.common import invalidInputError
from bigdl.llm.transformers.models.utils import use_flash_attention, use_esimd_sdp
from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, \ from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, \
append_kv_cache, is_enough_kv_cache_room_4_31 append_kv_cache, is_enough_kv_cache_room_4_31
from bigdl.llm.transformers.models.utils import init_fp8_kv_cache, append_fp8_kv_cache, \ from bigdl.llm.transformers.models.utils import init_fp8_kv_cache, append_fp8_kv_cache, \
@ -267,25 +269,43 @@ def baichuan_attention_forward_7b_origin(
past_key_value = (key_states, value_states) if use_cache else None past_key_value = (key_states, value_states) if use_cache else None
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) if not self.training and not hidden_states.requires_grad and \
use_flash_attention(query_states, key_states, attention_mask):
attn_output = F.scaled_dot_product_attention(query_states.to(device, dtype=torch.float16),
key_states.to(device, dtype=torch.float16),
value_states.to(device, dtype=torch.float16),
is_causal=True)
attn_weights = None
elif not self.training and not hidden_states.requires_grad and \
use_esimd_sdp(q_len, key_states.shape[2], self.head_dim, query_states):
import linear_fp16_esimd
attn_output = linear_fp16_esimd.sdp_forward(query_states,
key_states,
value_states)
attn_output = attn_output.view(query_states.shape)
attn_weights = None
else:
attn_weights = torch.matmul(query_states,
key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
invalidInputError(False, invalidInputError(False,
f"Attention weights should be of size " f"Attention weights should be of size "
f"{(bsz, self.num_heads, q_len, kv_seq_len)}" f"{(bsz, self.num_heads, q_len, kv_seq_len)}"
f", but is {attn_weights.size()}") f", but is {attn_weights.size()}")
if attention_mask is not None: if attention_mask is not None:
invalidInputError(attention_mask.size() == (bsz, 1, q_len, kv_seq_len), invalidInputError(attention_mask.size() == (bsz, 1, q_len, kv_seq_len),
f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, " f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, "
f"but is {attention_mask.size()}") f"but is {attention_mask.size()}")
attn_weights = attn_weights + attention_mask attn_weights = attn_weights + attention_mask
attn_weights = torch.max(attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min)) attn_weights = torch.max(attn_weights,
torch.tensor(torch.finfo(attn_weights.dtype).min))
# upcast attention to fp32 # upcast attention to fp32
attn_weights = nn.functional.softmax(attn_weights, dim=-1, attn_weights = nn.functional.softmax(attn_weights, dim=-1,
dtype=torch.float32).to(query_states.dtype) dtype=torch.float32).to(query_states.dtype)
attn_output = torch.matmul(attn_weights, value_states) attn_output = torch.matmul(attn_weights, value_states)
invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim), invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim),
f"`attn_output` should be of size " f"`attn_output` should be of size "
@ -300,7 +320,7 @@ def baichuan_attention_forward_7b_origin(
if not output_attentions: if not output_attentions:
attn_weights = None attn_weights = None
return attn_output, attn_weights, past_key_value return attn_output.to(hidden_states.dtype), attn_weights, past_key_value
def baichuan_attention_forward_13b( def baichuan_attention_forward_13b(
@ -502,4 +522,4 @@ def baichuan_attention_forward_13b_origin(
if not output_attentions: if not output_attentions:
attn_weights = None attn_weights = None
return attn_output, attn_weights, past_key_value return attn_output.to(hidden_states.dtype), attn_weights, past_key_value

37
python/llm/src/bigdl/llm/transformers/models/baichuan2.py
View file

@ -28,6 +28,7 @@ from bigdl.llm.transformers.models.utils import init_fp8_kv_cache, append_fp8_kv
restore_fp8_kv_cache, use_quantize_kv_cache restore_fp8_kv_cache, use_quantize_kv_cache
from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, \ from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, \
append_kv_cache, is_enough_kv_cache_room_4_31 append_kv_cache, is_enough_kv_cache_room_4_31
from bigdl.llm.transformers.models.utils import use_flash_attention, use_esimd_sdp
from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb, SILU from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb, SILU
from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb_no_cache_xpu from bigdl.llm.transformers.models.utils import apply_rotary_pos_emb_no_cache_xpu
from bigdl.llm.transformers.models.utils import mlp_fusion_check from bigdl.llm.transformers.models.utils import mlp_fusion_check
@ -271,16 +272,32 @@ def baichuan_attention_forward_7b_origin(
query_states, key_states, value_states, attn_bias=xops.LowerTriangularMask() query_states, key_states, value_states, attn_bias=xops.LowerTriangularMask()
) )
else: else:
if attention_mask is not None: if not self.training and not hidden_states.requires_grad and \
if attention_mask.dtype == torch.bool: use_flash_attention(query_states, key_states, attention_mask):
attention_mask.masked_fill_(attention_mask.logical_not(), float("-inf")) attn_output = F.scaled_dot_product_attention(query_states.to(dtype=torch.float16),
key_states.to(dtype=torch.float16),
value_states.to(dtype=torch.float16),
is_causal=True)
attn_weights = None
elif not self.training and not hidden_states.requires_grad and \
use_esimd_sdp(q_len, key_states.shape[2], self.head_dim, query_states):
import linear_fp16_esimd
attn_output = linear_fp16_esimd.sdp_forward(query_states,
key_states,
value_states)
attn_output = attn_output.view(query_states.shape)
attn_weights = None
else:
if attention_mask is not None:
if attention_mask.dtype == torch.bool:
attention_mask.masked_fill_(attention_mask.logical_not(), float("-inf"))
scaling_factor = 1 / math.sqrt(query_states.size(-1)) scaling_factor = 1 / math.sqrt(query_states.size(-1))
attn_output = torch.matmul(query_states * scaling_factor, key_states.transpose(-2, -1)) attn_output = torch.matmul(query_states * scaling_factor, key_states.transpose(-2, -1))
if attention_mask is not None: if attention_mask is not None:
attn_output += attention_mask attn_output += attention_mask
attn_output = torch.softmax(attn_output, -1) attn_output = torch.softmax(attn_output, -1)
attn_output = torch.matmul(attn_output, value_states) attn_output = torch.matmul(attn_output, value_states)
attn_output = attn_output.transpose(1, 2) attn_output = attn_output.transpose(1, 2)
attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim)
@ -289,7 +306,7 @@ def baichuan_attention_forward_7b_origin(
if not output_attentions: if not output_attentions:
attn_weights = None attn_weights = None
return attn_output, attn_weights, past_key_value return attn_output.to(hidden_states.dtype), attn_weights, past_key_value
def baichuan_attention_forward_13b( def baichuan_attention_forward_13b(

17
python/llm/src/bigdl/llm/transformers/models/qwen2.py
View file

@ -348,6 +348,13 @@ def qwen2_attention_forward_origin(
value_states = repeat_kv(value_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups)
if not self.training and not hidden_states.requires_grad and \ if not self.training and not hidden_states.requires_grad and \
use_flash_attention(query_states, key_states, attention_mask):
attn_output = F.scaled_dot_product_attention(query_states.to(device, dtype=torch.float16),
key_states.to(device, dtype=torch.float16),
value_states.to(device, dtype=torch.float16),
is_causal=True)
attn_weights = None
elif not self.training and not hidden_states.requires_grad and \
use_esimd_sdp(q_len, key_states.shape[2], self.head_dim, query_states): use_esimd_sdp(q_len, key_states.shape[2], self.head_dim, query_states):
import linear_fp16_esimd import linear_fp16_esimd
attn_output = linear_fp16_esimd.sdp_forward(query_states, attn_output = linear_fp16_esimd.sdp_forward(query_states,
@ -379,12 +386,12 @@ def qwen2_attention_forward_origin(
training=self.training) training=self.training)
attn_output = torch.matmul(attn_weights, value_states) attn_output = torch.matmul(attn_weights, value_states)
invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim), invalidInputError(attn_output.size() == (bsz, self.num_heads, q_len, self.head_dim),
"`attn_output` should be of size " "`attn_output` should be of size "
f"{(bsz, self.num_heads, q_len, self.head_dim)}," f"{(bsz, self.num_heads, q_len, self.head_dim)},"
f" but is {attn_output.size()}") f" but is {attn_output.size()}")
attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
attn_output = self.o_proj(attn_output) attn_output = self.o_proj(attn_output)