doing work to impliment model
This commit is contained in:
cm2435
parent
440ef5dfe3
commit
a1e2b0d43c
1 changed files with 290 additions and 0 deletions
290
unsloth/models/phi2.py
Normal file
290
unsloth/models/phi2.py
Normal file
|
|
@ -0,0 +1,290 @@
|
|||
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from .llama import *
|
||||
from ._utils import __version__
|
||||
from ..kernels.relu import relu_kernel
|
||||
|
||||
from transformers.models.phi.modeling_phi import (
|
||||
PhiAttention,
|
||||
PhiDecoderLayer,
|
||||
PhiModel,
|
||||
PhiForCausalLM,
|
||||
)
|
||||
# For Pytorch 2.1.1
|
||||
try:
|
||||
from transformers.models.phi.modeling_phi import (
|
||||
MistralFlashAttention2,
|
||||
)
|
||||
except:
|
||||
PhiFlashAttention2 = PhiAttention
|
||||
pass
|
||||
|
||||
|
||||
|
||||
def Phi2Attention_fast_forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
output_attentions: bool = False,
|
||||
use_cache: bool = False,
|
||||
padding_mask: Optional[torch.LongTensor] = None,
|
||||
*args, **kwargs,
|
||||
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
|
||||
|
||||
bsz, q_len, _ = hidden_states.size()
|
||||
Q, K, V = self.apply_qkv(self, hidden_states)
|
||||
|
||||
# Check for inference
|
||||
if use_cache and past_key_value is not None and q_len == 1:
|
||||
A, past_key_value = LlamaAttention_fast_forward_inference(
|
||||
self,
|
||||
hidden_states,
|
||||
past_key_value,
|
||||
position_ids,
|
||||
)
|
||||
return A, None, past_key_value
|
||||
pass
|
||||
#Get attention parameters
|
||||
n_heads = self.num_heads
|
||||
n_groups = self.num_key_value_groups
|
||||
n_kv_heads = self.num_key_value_heads
|
||||
head_dim = self.head_dim
|
||||
assert(n_kv_heads * n_groups == n_heads)
|
||||
|
||||
#.view() : (bsz, seq_len, embed_dim) -> (bsz, 1, n_attention_heads, head_dim)
|
||||
#transpose() : (bsz, 1, n_attention_heads, head_dim) -> (bsz, n_attention_heads, 1, head_dim)
|
||||
Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2)
|
||||
K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)
|
||||
|
||||
kv_seq_len = K.shape[-2]
|
||||
if past_key_value is not None:
|
||||
kv_seq_len += past_key_value[0].shape[-2]
|
||||
|
||||
if position_ids is None:
|
||||
cos = self.rotary_emb.cos_cached
|
||||
sin = self.rotary_emb.sin_cached
|
||||
Q, K = fast_rope_embedding(Q, K, cos, sin)
|
||||
else:
|
||||
cos, sin = self.rotary_emb(V, seq_len = kv_seq_len)
|
||||
Q, K = inplace_rope_embedding(Q, K, cos, sin, position_ids)
|
||||
pass
|
||||
|
||||
if past_key_value is not None:
|
||||
# reuse k, v, self_attention
|
||||
K = torch.cat([past_key_value[0], K], dim = 2)
|
||||
V = torch.cat([past_key_value[1], V], dim = 2)
|
||||
past_key_value = (K, V) if use_cache else None
|
||||
|
||||
|
||||
# Attention module
|
||||
if (not HAS_FLASH_ATTENTION):
|
||||
# Xformers memory efficient attention
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
M = bsz * q_len
|
||||
|
||||
has_swa = isinstance(causal_mask, xformers.attn_bias.BlockDiagonalCausalMask)
|
||||
|
||||
# Group query attention
|
||||
K = K .view(bsz, q_len, n_kv_heads, 1, head_dim)
|
||||
V = V .view(bsz, q_len, n_kv_heads, 1, head_dim)
|
||||
K = K.expand(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
V = V.expand(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
if hidden_states.requires_grad:
|
||||
K = K.reshape(bsz, q_len, n_heads, head_dim)
|
||||
V = V.reshape(bsz, q_len, n_heads, head_dim)
|
||||
|
||||
if has_swa:
|
||||
Q = Q.view(1, M, n_heads, head_dim)
|
||||
K = K.view(1, M, n_heads, head_dim)
|
||||
V = V.view(1, M, n_heads, head_dim)
|
||||
pass
|
||||
else:
|
||||
# Xformers does support the forward pass though
|
||||
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
|
||||
|
||||
if has_swa:
|
||||
Q = Q.view(1, M, n_kv_heads, n_groups, head_dim)
|
||||
K = K.view(1, M, n_kv_heads, n_groups, head_dim)
|
||||
V = V.view(1, M, n_kv_heads, n_groups, head_dim)
|
||||
pass
|
||||
pass
|
||||
|
||||
elif HAS_FLASH_ATTENTION:
|
||||
Q = Q.transpose(1, 2)
|
||||
K = K.transpose(1, 2)
|
||||
V = V.transpose(1, 2)
|
||||
sw = getattr(self.config, "sliding_window")
|
||||
sw = q_len if sw is None else sw
|
||||
window = (-1, -1) if (q_len <= sw) else (sw, sw)
|
||||
A = flash_attn_func(Q, K, V, causal = True, window_size = window)
|
||||
else:
|
||||
# Grouped query attention
|
||||
# if n_groups != 1:
|
||||
K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim)
|
||||
V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim)
|
||||
K = K.reshape(bsz, n_heads, q_len, head_dim)
|
||||
V = V.reshape(bsz, n_heads, q_len, head_dim)
|
||||
# Needs (batch_size, n_heads, seq_len, head_dim)
|
||||
# is_casual and attention_mask must not be both set!
|
||||
A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False)
|
||||
# Go back to (batch_size, seq_len, n_heads, head_dim)
|
||||
A = A.transpose(1, 2)
|
||||
pass
|
||||
|
||||
attn_output = A.reshape(bsz, q_len, self.hidden_size)
|
||||
attn_output = self.apply_o(self, attn_output)
|
||||
attn_weights = None
|
||||
return attn_output, attn_weights, past_key_value
|
||||
pass
|
||||
|
||||
inplace_rope_embedding
|
||||
def fast_mlp_inference(self, X):
|
||||
gate = self.gate_proj(X)
|
||||
up = self.up_proj(X)
|
||||
gate = relu_kernel(gate, inplace = True)
|
||||
gate *= up
|
||||
X = self.down_proj(gate)
|
||||
return X
|
||||
pass
|
||||
|
||||
class FastPhi2Model(FastLlamaModel):
|
||||
|
||||
@staticmethod
|
||||
def pre_patch():
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
def from_pretrained(
|
||||
model_name = "unsloth/llama-2-7b-bnb-4bit", #TODO: update me.
|
||||
max_seq_length = 4096,
|
||||
dtype = None,
|
||||
load_in_4bit = True,
|
||||
token = None,
|
||||
device_map = "sequential",
|
||||
rope_scaling = None,
|
||||
fix_tokenizer = True,
|
||||
):
|
||||
|
||||
SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported()
|
||||
gpu_stats = torch.cuda.get_device_properties(0)
|
||||
max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
|
||||
|
||||
statistics = \
|
||||
f"==((====))== Unsloth: Fast Mistral patching release {__version__}\n"\
|
||||
f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB\n"\
|
||||
f"O^O/ \_/ \\ CUDA capability = {gpu_stats.major}.{gpu_stats.minor}. Xformers = {xformers_version}. FA = {HAS_FLASH_ATTENTION}.\n"\
|
||||
f"\ / Pytorch version: {torch.__version__}. CUDA Toolkit = {torch.version.cuda}\n"\
|
||||
f' "-____-" bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. Platform = {platform_system}\n'
|
||||
logger.warning_once(statistics)
|
||||
FastPhi2Model.pre_patch()
|
||||
|
||||
if dtype is None:
|
||||
dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16
|
||||
elif dtype == torch.bfloat16 and not SUPPORTS_BFLOAT16:
|
||||
logger.warning_once("Device does not support bfloat16. Will change to float16.")
|
||||
dtype = torch.float16
|
||||
|
||||
assert(dtype == torch.float16 or dtype == torch.bfloat16 or dtype == torch.float32)
|
||||
|
||||
# RoPE scaling
|
||||
model_max_seq_length = \
|
||||
AutoConfig.from_pretrained(model_name, token = token).max_position_embeddings
|
||||
|
||||
if (rope_scaling is None) and (max_seq_length > model_max_seq_length):
|
||||
rope_scaling = max_seq_length / model_max_seq_length
|
||||
logger.warning_once(
|
||||
f"Unsloth: {model_name} can only handle sequence lengths of at most "\
|
||||
f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "\
|
||||
f"{round(rope_scaling, 3)}, it can be magically be extended to "\
|
||||
f"{max_seq_length}!"
|
||||
)
|
||||
rope_scaling = {"type": "linear", "factor": rope_scaling,}
|
||||
pass
|
||||
|
||||
bnb_config = None
|
||||
if load_in_4bit:
|
||||
bnb_config = BitsAndBytesConfig(
|
||||
load_in_4bit = True,
|
||||
bnb_4bit_use_double_quant = True,
|
||||
bnb_4bit_quant_type = "nf4",
|
||||
bnb_4bit_compute_dtype = dtype,
|
||||
)
|
||||
|
||||
# https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/discussions/12
|
||||
# RoPE Scaling's max_position_embeddings must be updated
|
||||
max_position_embeddings = max(max_seq_length, model_max_seq_length)
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
device_map = device_map,
|
||||
torch_dtype = dtype,
|
||||
quantization_config = bnb_config,
|
||||
token = token,
|
||||
# rope_scaling = rope_scaling,
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_name,
|
||||
model_max_length = max_seq_length,
|
||||
padding_side = "right",
|
||||
token = token,
|
||||
)
|
||||
|
||||
model, tokenizer = patch_tokenizer(model, tokenizer)
|
||||
model = FastPhi2Model.post_patch(model)
|
||||
|
||||
# Patch up QKV / O and MLP
|
||||
for idx, layer in enumerate(model.model.layers):
|
||||
layer.self_attn.apply_qkv = original_apply_qkv
|
||||
layer.self_attn.apply_o = original_apply_o
|
||||
pass
|
||||
|
||||
# Save max_seq_length
|
||||
model.max_seq_length = max_position_embeddings
|
||||
internal_model = model
|
||||
while hasattr(internal_model, "model"):
|
||||
internal_model.max_seq_length = max_position_embeddings
|
||||
internal_model = internal_model.model
|
||||
pass
|
||||
internal_model.max_seq_length = max_position_embeddings
|
||||
|
||||
# We check the tokenizer first for errors
|
||||
if fix_tokenizer:
|
||||
tokenizer = check_tokenizer(
|
||||
model = model,
|
||||
tokenizer = tokenizer,
|
||||
model_name = model_name,
|
||||
model_max_length = max_seq_length,
|
||||
padding_side = "right",
|
||||
token = token,
|
||||
)
|
||||
pass
|
||||
patch_saving_functions(tokenizer)
|
||||
# Fix up config for transformers uploading PEFT
|
||||
name = model.config._name_or_path
|
||||
if name.startswith("unsloth/") and name.endswith("-bnb-4bit"):
|
||||
name = name[:len(name) - len("-bnb-4bit")]
|
||||
model.config.update({"_name_or_path" : name})
|
||||
pass
|
||||
|
||||
# Log Unsloth version for future fastpaths for inference
|
||||
model.config.update({"unsloth_version" : __version__})
|
||||
|
||||
return model, tokenizer
|
||||
pass
|
||||
Loading…
Add table
Add a link
Reference in a new issue