Refine model of qwen.
This commit is contained in:
Colin
parent
40ae899515
commit
7c047f0b32
44
qwen/demo.py
44
qwen/demo.py
|
|
@ -24,33 +24,6 @@ model = QWenLMHeadModel(config)
|
|||
|
||||
print(model)
|
||||
|
||||
# QWenLMHeadModel(
|
||||
# (transformer): QWenModel(
|
||||
# (wte): Embedding(151936, 2048)
|
||||
# (drop): Dropout(p=0.0, inplace=False)
|
||||
# (rotary_emb): RotaryEmbedding()
|
||||
# (h): ModuleList(
|
||||
# (0-23): 24 x QWenBlock(
|
||||
# (ln_1): RMSNorm()
|
||||
# (attn): QWenAttention(
|
||||
# (c_attn): Linear(in_features=2048, out_features=6144, bias=True)
|
||||
# (c_proj): Linear(in_features=2048, out_features=2048, bias=False)
|
||||
# (attn_dropout): Dropout(p=0.0, inplace=False)
|
||||
# )
|
||||
# (ln_2): RMSNorm()
|
||||
# (mlp): QWenMLP(
|
||||
# (w1): Linear(in_features=2048, out_features=5504, bias=False)
|
||||
# (w2): Linear(in_features=2048, out_features=5504, bias=False)
|
||||
# (c_proj): Linear(in_features=5504, out_features=2048, bias=False)
|
||||
# )
|
||||
# )
|
||||
# )
|
||||
# (ln_f): RMSNorm()
|
||||
# )
|
||||
# (lm_head): Linear(in_features=2048, out_features=151936, bias=False)
|
||||
# )
|
||||
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
|
||||
model = model.from_pretrained(model_dir).cuda()
|
||||
|
||||
|
|
@ -72,22 +45,9 @@ print(decode_tokens)
|
|||
# <|im_start|>assistant
|
||||
# 日本的首都东京。<|im_end|><|endoftext|>
|
||||
|
||||
|
||||
# # 第一轮对话
|
||||
# response, history, decode_tokens = model.chat(tokenizer, "你好", "", history=None)
|
||||
# print(decode_tokens)
|
||||
# # 你好！很高兴为你提供帮助。
|
||||
|
||||
# 第二轮对话
|
||||
response, history, decode_tokens = model.chat(tokenizer, "给我讲一个年轻人奋斗创业最终取得成功的故事。", "", history=None)
|
||||
print(decode_tokens)
|
||||
|
||||
|
||||
# <|im_start|>system
|
||||
# You are a helpful assistant.<|im_end|>
|
||||
# <|im_start|>user
|
||||
# 你好<|im_end|>
|
||||
# <|im_start|>assistant
|
||||
# 莎士比亚是头一个使用“你好”这个词的文学家，他在《哈姆雷特》中写道：“你是谁？你在哪儿？
|
||||
# ”他的这一段话，通常被认为是最早的使用“你好”这个词的文学记载。这句话在英国语中非常常见，
|
||||
# 特别是在正式或礼貌的情况下。<|im_end|><|endoftext|>
|
||||
if decode_tokens.split("\n")[-2] != """这个故事告诉我们，只要我们有决心和毅力，就一定能够克服困难，实现我们的梦想。<|im_end|>""":
|
||||
raise ()
|
||||
|
|
|
|||
|
|
@ -1,47 +1,29 @@
|
|||
import copy
|
||||
import math
|
||||
import inspect
|
||||
import os
|
||||
import sys
|
||||
import gc
|
||||
from tqdm import auto as tqdm_lib
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Optional, Tuple, Union, Callable, List, Any, Generator
|
||||
from typing import Optional, Tuple, Union, Callable, List, Any, Generator
|
||||
from einops import rearrange
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.checkpoint
|
||||
|
||||
from torch.nn import CrossEntropyLoss
|
||||
from transformers import PreTrainedTokenizer, GenerationConfig, StoppingCriteriaList
|
||||
from transformers.generation.logits_process import LogitsProcessorList
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from transformers.generation.streamers import BaseStreamer
|
||||
from transformers.generation.utils import GenerateOutput
|
||||
from transformers.modeling_outputs import (
|
||||
BaseModelOutputWithPast,
|
||||
CausalLMOutputWithPast,
|
||||
)
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
from transformers.utils import logging
|
||||
|
||||
from torch import nn
|
||||
from einops import rearrange
|
||||
from safetensors.torch import load_file as safe_load_file
|
||||
from safetensors.torch import save_file as safe_save_file
|
||||
|
||||
from transformers.generation.utils import GenerateOutput
|
||||
from configuration_qwen import QWenConfig
|
||||
from qwen_generation_utils import (
|
||||
HistoryType,
|
||||
make_context,
|
||||
decode_tokens,
|
||||
StopWordsLogitsProcessor,
|
||||
)
|
||||
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import load_file as safe_load_file
|
||||
from safetensors.torch import save_file as safe_save_file
|
||||
|
||||
import sys
|
||||
|
||||
sys.path.append("..")
|
||||
from tools import show
|
||||
from tools import mem_tracker
|
||||
|
|
@ -50,39 +32,30 @@ from tools import mem_tracker
|
|||
# tracker.track()
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def _norm(self, x):
|
||||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
|
||||
def forward(self, x):
|
||||
return self._norm(x.float()).type_as(x) * self.weight
|
||||
|
||||
|
||||
class QWenAttention(nn.Module):
|
||||
def __init__(self, config, index):
|
||||
super().__init__()
|
||||
|
||||
self.register_buffer("masked_bias", torch.tensor(-1e4), persistent=False)
|
||||
self.seq_length = config.seq_length
|
||||
|
||||
self.hidden_size = config.hidden_size
|
||||
self.split_size = config.hidden_size
|
||||
self.num_heads = config.num_attention_heads
|
||||
self.head_dim = self.hidden_size // self.num_heads
|
||||
|
||||
self.scale_attn_weights = True
|
||||
|
||||
self.projection_size = config.kv_channels * config.num_attention_heads
|
||||
|
||||
assert self.projection_size % config.num_attention_heads == 0
|
||||
self.hidden_size_per_attention_head = self.projection_size // config.num_attention_heads
|
||||
|
||||
self.c_attn = nn.Linear(config.hidden_size, 3 * self.projection_size)
|
||||
self.c_proj = nn.Linear(config.hidden_size, self.projection_size, bias=not config.no_bias)
|
||||
|
||||
self.use_dynamic_ntk = config.use_dynamic_ntk
|
||||
|
||||
logn_list = [math.log(i, self.seq_length) if i > self.seq_length else 1 for i in range(1, 32768)]
|
||||
logn_tensor = torch.tensor(logn_list)[None, :, None, None]
|
||||
self.register_buffer("logn_tensor", logn_tensor, persistent=False)
|
||||
|
||||
self.attn_dropout = nn.Dropout(config.attn_dropout_prob)
|
||||
self.softmax_in_fp32 = config.softmax_in_fp32 if hasattr(config, "softmax_in_fp32") else False
|
||||
cache_dtype = torch.float
|
||||
self.cache_qmax = torch.tensor(torch.iinfo(torch.uint8).max, dtype=cache_dtype)
|
||||
self.cache_qmin = torch.tensor(torch.iinfo(torch.uint8).min, dtype=cache_dtype)
|
||||
self.index = index
|
||||
|
||||
def _split_heads(self, tensor, num_heads, attn_head_size):
|
||||
|
|
@ -95,53 +68,6 @@ class QWenAttention(nn.Module):
|
|||
new_shape = tensor.size()[:-2] + (num_heads * attn_head_size,)
|
||||
return tensor.view(new_shape)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: Optional[Tuple[torch.FloatTensor]],
|
||||
rotary_pos_emb_list: Optional[List[List[torch.Tensor]]] = None,
|
||||
):
|
||||
mixed_x_layer = self.c_attn(hidden_states)
|
||||
query, key, value = mixed_x_layer.split(self.split_size, dim=2)
|
||||
query = self._split_heads(query, self.num_heads, self.head_dim)
|
||||
key = self._split_heads(key, self.num_heads, self.head_dim)
|
||||
value = self._split_heads(value, self.num_heads, self.head_dim)
|
||||
|
||||
rotary_pos_emb = rotary_pos_emb_list[0]
|
||||
rotary_pos_emb = [i[:, -query.shape[1] :, :, :] for i in rotary_pos_emb]
|
||||
rotary_pos_emb = (rotary_pos_emb,) * 2
|
||||
q_pos_emb, k_pos_emb = rotary_pos_emb
|
||||
# Slice the pos emb for current inference
|
||||
query = apply_rotary_pos_emb(query, q_pos_emb)
|
||||
key = apply_rotary_pos_emb(key, k_pos_emb)
|
||||
|
||||
key_size = key.size(1)
|
||||
if key_size > self.seq_length and not self.training:
|
||||
seq_start = key.size(1) - query.size(1)
|
||||
seq_end = key.size(1)
|
||||
logn_tensor = self.logn_tensor[:, seq_start:seq_end, :, :].type_as(query)
|
||||
query = query * logn_tensor.expand_as(query)
|
||||
|
||||
key_size = key.size(1)
|
||||
causal_mask = torch.tril(torch.ones((key_size, key_size), dtype=torch.bool, device=query.device)).view(
|
||||
1, 1, key_size, key_size
|
||||
)
|
||||
query = query.permute(0, 2, 1, 3)
|
||||
key = key.permute(0, 2, 1, 3)
|
||||
value = value.permute(0, 2, 1, 3)
|
||||
|
||||
# qk = query @ key.transpose(-2, -1)
|
||||
# qk = qk[0]
|
||||
# prePath = "../generated/query_matmul_key/img/"
|
||||
# show.DumpTensorToImage(
|
||||
# qk, prePath + "q_matmul_k_sequence_" + str(key_size) + "_layer_" + str(self.index) + ".png"
|
||||
# )
|
||||
|
||||
attn_output = F.scaled_dot_product_attention(query, key, value, attn_mask=causal_mask).transpose(1, 2)
|
||||
context_layer = self._merge_heads(attn_output, self.num_heads, self.head_dim)
|
||||
attn_output = self.c_proj(context_layer)
|
||||
|
||||
return attn_output
|
||||
|
||||
|
||||
class QWenMLP(nn.Module):
|
||||
def __init__(self, config):
|
||||
|
|
@ -151,110 +77,60 @@ class QWenMLP(nn.Module):
|
|||
self.w2 = nn.Linear(config.hidden_size, ff_dim_in, bias=not config.no_bias)
|
||||
self.c_proj = nn.Linear(ff_dim_in, config.hidden_size, bias=not config.no_bias)
|
||||
|
||||
def forward(self, hidden_states):
|
||||
a1 = self.w1(hidden_states)
|
||||
a2 = self.w2(hidden_states)
|
||||
intermediate_parallel = a1 * F.silu(a2)
|
||||
output = self.c_proj(intermediate_parallel)
|
||||
return output
|
||||
|
||||
|
||||
class QWenBlock(nn.Module):
|
||||
def __init__(self, config, index):
|
||||
super().__init__()
|
||||
hidden_size = config.hidden_size
|
||||
|
||||
self.ln_1 = RMSNorm(
|
||||
hidden_size,
|
||||
config.hidden_size,
|
||||
eps=config.layer_norm_epsilon,
|
||||
)
|
||||
self.attn = QWenAttention(config, index)
|
||||
self.ln_2 = RMSNorm(
|
||||
hidden_size,
|
||||
config.hidden_size,
|
||||
eps=config.layer_norm_epsilon,
|
||||
)
|
||||
self.mlp = QWenMLP(config)
|
||||
self.index = index
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: Optional[Tuple[torch.FloatTensor]],
|
||||
rotary_pos_emb_list: Optional[List[List[torch.Tensor]]] = None,
|
||||
):
|
||||
layernorm_output = self.ln_1(hidden_states)
|
||||
|
||||
attn_outputs = self.attn(layernorm_output, rotary_pos_emb_list)
|
||||
attn_output = attn_outputs[0]
|
||||
residual = hidden_states
|
||||
layernorm_input = attn_output + residual
|
||||
|
||||
layernorm_output = self.ln_2(layernorm_input)
|
||||
residual = layernorm_input
|
||||
mlp_output = self.mlp(layernorm_output)
|
||||
hidden_states = residual + mlp_output
|
||||
return hidden_states
|
||||
|
||||
|
||||
class QWenPreTrainedModel(nn.Module):
|
||||
config_class = QWenConfig
|
||||
base_model_prefix = "transformer"
|
||||
is_parallelizable = False
|
||||
supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["QWenBlock"]
|
||||
|
||||
def __init__(self, *inputs, **kwargs):
|
||||
super().__init__()
|
||||
|
||||
|
||||
class QWenModel(QWenPreTrainedModel):
|
||||
class QWenModel(nn.Module):
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.num_hidden_layers = config.num_hidden_layers
|
||||
self.embed_dim = config.hidden_size
|
||||
|
||||
self.use_dynamic_ntk = config.use_dynamic_ntk
|
||||
self.seq_length = config.seq_length
|
||||
|
||||
self.wte = nn.Embedding(self.vocab_size, self.embed_dim)
|
||||
|
||||
super().__init__()
|
||||
self.wte = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.drop = nn.Dropout(config.emb_dropout_prob)
|
||||
|
||||
if config.rotary_pct == 1.0:
|
||||
self.rotary_ndims = None
|
||||
else:
|
||||
assert config.rotary_pct < 1
|
||||
self.rotary_ndims = int(config.kv_channels * config.rotary_pct)
|
||||
dim = self.rotary_ndims if self.rotary_ndims is not None else config.kv_channels
|
||||
self.rotary_emb = RotaryEmbedding(dim, base=config.rotary_emb_base)
|
||||
dim = config.kv_channels
|
||||
|
||||
self.h = nn.ModuleList([QWenBlock(config, i) for i in range(config.num_hidden_layers)])
|
||||
self.ln_f = RMSNorm(
|
||||
self.embed_dim,
|
||||
config.hidden_size,
|
||||
eps=config.layer_norm_epsilon,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
):
|
||||
input_shape = input_ids.size()
|
||||
input_ids = input_ids.view(-1, input_shape[-1])
|
||||
batch_size = input_ids.shape[0]
|
||||
hidden_states = self.wte(input_ids)
|
||||
kv_seq_len = hidden_states.size()[1]
|
||||
rotary_pos_emb_list = [self.rotary_emb(kv_seq_len, ntk_alpha=1.0)]
|
||||
self.dim = dim
|
||||
self.base = config.rotary_emb_base
|
||||
inv_freq = 1.0 / (self.base ** (torch.arange(0, dim, 2).float() / dim))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
self._rotary_pos_emb_cache = None
|
||||
self._seq_len_cached = 0
|
||||
self._ntk_alpha_cached = 1.0
|
||||
|
||||
hidden_states = self.drop(hidden_states)
|
||||
output_shape = input_shape + (hidden_states.size(-1),)
|
||||
def update_rotary_pos_emb_cache(self, seqlen, ntk_alpha=1.0):
|
||||
if seqlen > self._seq_len_cached or ntk_alpha != self._ntk_alpha_cached:
|
||||
base = self.base * ntk_alpha ** (self.dim / (self.dim - 2))
|
||||
self.inv_freq = 1.0 / (
|
||||
base ** (torch.arange(0, self.dim, 2, device=self.inv_freq.device).float() / self.dim)
|
||||
)
|
||||
self._seq_len_cached = max(2 * seqlen, 16)
|
||||
self._ntk_alpha_cached = ntk_alpha
|
||||
seq = torch.arange(self._seq_len_cached, device=self.inv_freq.device)
|
||||
freqs = torch.outer(seq.type_as(self.inv_freq), self.inv_freq)
|
||||
|
||||
all_hidden_states = None
|
||||
for block in self.h:
|
||||
hidden_states = block(hidden_states, rotary_pos_emb_list=rotary_pos_emb_list)
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
emb = rearrange(emb, "n d -> 1 n 1 d")
|
||||
|
||||
hidden_states = self.ln_f(hidden_states)
|
||||
hidden_states = hidden_states.view(output_shape)
|
||||
return BaseModelOutputWithPast(last_hidden_state=hidden_states, hidden_states=all_hidden_states)
|
||||
cos, sin = emb.cos(), emb.sin()
|
||||
self._rotary_pos_emb_cache = [cos, sin]
|
||||
|
||||
|
||||
class QWenLMHeadModel(nn.Module):
|
||||
|
|
@ -264,51 +140,8 @@ class QWenLMHeadModel(nn.Module):
|
|||
|
||||
self.transformer = QWenModel(config)
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
self.generation_config = GenerationConfig.from_model_config(config)
|
||||
|
||||
def prepare_inputs_for_generation(self, input_ids, **kwargs):
|
||||
model_inputs = {"input_ids": input_ids}
|
||||
return model_inputs
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
transformer_outputs = self.transformer(
|
||||
input_ids,
|
||||
)
|
||||
hidden_states = transformer_outputs[0]
|
||||
|
||||
lm_logits = self.lm_head(hidden_states)
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
labels = labels.to(lm_logits.device)
|
||||
shift_logits = lm_logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
loss_fct = CrossEntropyLoss()
|
||||
loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
|
||||
|
||||
# shift_labels = torch.ones([1,19]).to(lm_logits.device).to(torch.int64)
|
||||
# shift_logits = lm_logits[..., :-1, :].contiguous()
|
||||
# loss_fct = CrossEntropyLoss()
|
||||
# loss = loss_fct(
|
||||
# shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
|
||||
# )
|
||||
# loss.backward()
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=lm_logits,
|
||||
hidden_states=transformer_outputs.hidden_states,
|
||||
attentions=transformer_outputs.attentions,
|
||||
)
|
||||
|
||||
def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]]):
|
||||
load_in_8bit = False
|
||||
load_in_4bit = False
|
||||
|
||||
pretrained_model_name_or_path = str(pretrained_model_name_or_path)
|
||||
resolved_archive_file = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json")
|
||||
print(f"loading weights file {resolved_archive_file}")
|
||||
|
|
@ -317,8 +150,6 @@ class QWenLMHeadModel(nn.Module):
|
|||
shard_filenames = sorted(set(index["weight_map"].values()))
|
||||
resolved_archive_file = [os.path.join(pretrained_model_name_or_path, f) for f in shard_filenames]
|
||||
model = cls._load_pretrained_model(resolved_archive_file)
|
||||
model.is_loaded_in_4bit = load_in_4bit
|
||||
model.is_loaded_in_8bit = load_in_8bit
|
||||
return model
|
||||
|
||||
def _load_state_dict_into_model(self, model_to_load, state_dict, start_prefix):
|
||||
|
|
@ -358,29 +189,22 @@ class QWenLMHeadModel(nn.Module):
|
|||
@torch.no_grad()
|
||||
def chat(
|
||||
self,
|
||||
tokenizer: PreTrainedTokenizer,
|
||||
tokenizer,
|
||||
query: str,
|
||||
query_assistant: str,
|
||||
history: Optional[HistoryType],
|
||||
system: str = "You are a helpful assistant.",
|
||||
**kwargs,
|
||||
) -> Tuple[str, HistoryType]:
|
||||
generation_config = self.generation_config
|
||||
|
||||
if history is None:
|
||||
history = []
|
||||
else:
|
||||
history = copy.deepcopy(history)
|
||||
|
||||
stop_words_ids = []
|
||||
|
||||
raw_text, context_tokens = make_context(tokenizer, query, query_assistant, history=history, system=system)
|
||||
|
||||
stop_words_ids.extend([[tokenizer.im_end_id], [tokenizer.im_start_id]])
|
||||
input_ids = torch.tensor([context_tokens]).to(next(self.parameters()).device)
|
||||
outputs = self.generate(
|
||||
input_ids,
|
||||
stop_words_ids=stop_words_ids,
|
||||
tokenizer=tokenizer,
|
||||
**kwargs,
|
||||
)
|
||||
|
|
@ -397,31 +221,20 @@ class QWenLMHeadModel(nn.Module):
|
|||
def generate(
|
||||
self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
stop_words_ids=[],
|
||||
tokenizer=None,
|
||||
prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
|
||||
**kwargs,
|
||||
) -> Union[GenerateOutput, torch.LongTensor]:
|
||||
generation_config = self.generation_config
|
||||
generation_config = copy.deepcopy(generation_config)
|
||||
model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs
|
||||
generation_config.validate()
|
||||
pad_token_id = self.config.pad_token_id
|
||||
eos_token_id_tensor = torch.tensor([self.config.eos_token_id]).to(input_ids.device)
|
||||
|
||||
pad_token_id = generation_config.pad_token_id
|
||||
eos_token_id_tensor = torch.tensor([generation_config.eos_token_id]).to(input_ids.device)
|
||||
|
||||
scores = None
|
||||
# keep track of which sequences are already finished
|
||||
unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
|
||||
|
||||
this_peer_finished = False
|
||||
# auto-regressive generation
|
||||
while True:
|
||||
model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
|
||||
|
||||
# forward pass to get next token
|
||||
outputs = self(**model_inputs)
|
||||
next_token_scores = outputs.logits[:, -1, :]
|
||||
outputs = forwardQWen(self, input_ids)
|
||||
next_token_scores = outputs[:, -1, :]
|
||||
|
||||
# repetition_penalty
|
||||
penalty = self.config.repetition_penalty
|
||||
|
|
@ -475,64 +288,120 @@ class QWenLMHeadModel(nn.Module):
|
|||
return input_ids
|
||||
|
||||
|
||||
class RotaryEmbedding(torch.nn.Module):
|
||||
def __init__(self, dim, base=10000):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.base = base
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
self._rotary_pos_emb_cache = None
|
||||
self._seq_len_cached = 0
|
||||
self._ntk_alpha_cached = 1.0
|
||||
def forwardAttention(
|
||||
attention,
|
||||
hidden_states: Optional[Tuple[torch.FloatTensor]],
|
||||
rotary_pos_emb_list: Optional[List[List[torch.Tensor]]] = None,
|
||||
):
|
||||
def apply_rotary_pos_emb(t, freqs):
|
||||
def _rotate_half(x):
|
||||
x = rearrange(x, "... (j d) -> ... j d", j=2)
|
||||
x1, x2 = x.unbind(dim=-2)
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
def update_rotary_pos_emb_cache(self, seqlen, ntk_alpha=1.0):
|
||||
if seqlen > self._seq_len_cached or ntk_alpha != self._ntk_alpha_cached:
|
||||
base = self.base * ntk_alpha ** (self.dim / (self.dim - 2))
|
||||
self.inv_freq = 1.0 / (
|
||||
base ** (torch.arange(0, self.dim, 2, device=self.inv_freq.device).float() / self.dim)
|
||||
)
|
||||
self._seq_len_cached = max(2 * seqlen, 16)
|
||||
self._ntk_alpha_cached = ntk_alpha
|
||||
seq = torch.arange(self._seq_len_cached, device=self.inv_freq.device)
|
||||
freqs = torch.outer(seq.type_as(self.inv_freq), self.inv_freq)
|
||||
rot_dim = freqs[0].shape[-1]
|
||||
cos, sin = freqs
|
||||
t_float = t.float()
|
||||
t_rot, t_pass = t_float[..., :rot_dim], t_float[..., rot_dim:]
|
||||
t_rot = (t_rot * cos) + (_rotate_half(t_rot) * sin)
|
||||
return torch.cat((t_rot, t_pass), dim=-1).type_as(t)
|
||||
|
||||
emb = torch.cat((freqs, freqs), dim=-1)
|
||||
emb = rearrange(emb, "n d -> 1 n 1 d")
|
||||
atten = attention
|
||||
mixed_x_layer = atten.c_attn(hidden_states)
|
||||
query, key, value = mixed_x_layer.split(atten.split_size, dim=2)
|
||||
query = atten._split_heads(query, atten.num_heads, atten.head_dim)
|
||||
key = atten._split_heads(key, atten.num_heads, atten.head_dim)
|
||||
value = atten._split_heads(value, atten.num_heads, atten.head_dim)
|
||||
|
||||
cos, sin = emb.cos(), emb.sin()
|
||||
self._rotary_pos_emb_cache = [cos, sin]
|
||||
rotary_pos_emb = rotary_pos_emb_list[0]
|
||||
rotary_pos_emb = [i[:, -query.shape[1] :, :, :] for i in rotary_pos_emb]
|
||||
rotary_pos_emb = (rotary_pos_emb,) * 2
|
||||
query = apply_rotary_pos_emb(query, rotary_pos_emb[0])
|
||||
key = apply_rotary_pos_emb(key, rotary_pos_emb[1])
|
||||
|
||||
def forward(self, max_seq_len, ntk_alpha=1.0):
|
||||
self.update_rotary_pos_emb_cache(max_seq_len, ntk_alpha)
|
||||
cos, sin = self._rotary_pos_emb_cache
|
||||
return [cos[:, :max_seq_len], sin[:, :max_seq_len]]
|
||||
key_size = key.size(1)
|
||||
causal_mask = torch.tril(torch.ones((key_size, key_size), dtype=torch.bool, device=query.device)).view(
|
||||
1, 1, key_size, key_size
|
||||
)
|
||||
query = query.permute(0, 2, 1, 3)
|
||||
key = key.permute(0, 2, 1, 3)
|
||||
value = value.permute(0, 2, 1, 3)
|
||||
|
||||
# qk = query @ key.transpose(-2, -1)
|
||||
# qk = qk[0]
|
||||
# prePath = "../generated/query_matmul_key/img/"
|
||||
# show.DumpTensorToImage(
|
||||
# qk, prePath + "q_matmul_k_sequence_" + str(key_size) + "_layer_" + str(self.index) + ".png"
|
||||
# )
|
||||
|
||||
attn_output = F.scaled_dot_product_attention(query, key, value, attn_mask=causal_mask).transpose(1, 2)
|
||||
context_layer = atten._merge_heads(attn_output, atten.num_heads, atten.head_dim)
|
||||
attn_output = atten.c_proj(context_layer)
|
||||
|
||||
return attn_output
|
||||
|
||||
|
||||
def _rotate_half(x):
|
||||
x = rearrange(x, "... (j d) -> ... j d", j=2)
|
||||
x1, x2 = x.unbind(dim=-2)
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
def forwardQWenBlock(
|
||||
block,
|
||||
hidden_states: Optional[Tuple[torch.FloatTensor]],
|
||||
rotary_pos_emb_list: Optional[List[List[torch.Tensor]]] = None,
|
||||
):
|
||||
layernorm_output = block.ln_1(hidden_states)
|
||||
|
||||
attn_outputs = forwardAttention(block.attn, layernorm_output, rotary_pos_emb_list)
|
||||
attn_output = attn_outputs[0]
|
||||
layernorm_input = attn_output + hidden_states
|
||||
|
||||
layernorm_output = block.ln_2(layernorm_input)
|
||||
a1 = block.mlp.w1(layernorm_output)
|
||||
a2 = block.mlp.w2(layernorm_output)
|
||||
intermediate_parallel = a1 * F.silu(a2)
|
||||
mlp_output = block.mlp.c_proj(intermediate_parallel)
|
||||
|
||||
hidden_states = layernorm_input + mlp_output
|
||||
return hidden_states
|
||||
|
||||
|
||||
def apply_rotary_pos_emb(t, freqs):
|
||||
rot_dim = freqs[0].shape[-1]
|
||||
cos, sin = freqs
|
||||
t_float = t.float()
|
||||
t_rot, t_pass = t_float[..., :rot_dim], t_float[..., rot_dim:]
|
||||
t_rot = (t_rot * cos) + (_rotate_half(t_rot) * sin)
|
||||
return torch.cat((t_rot, t_pass), dim=-1).type_as(t)
|
||||
def forwardQWen(
|
||||
qwen,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
):
|
||||
transfm = qwen.transformer
|
||||
input_shape = input_ids.size()
|
||||
input_ids = input_ids.view(-1, input_shape[-1])
|
||||
hidden_states = transfm.wte(input_ids)
|
||||
kv_seq_len = hidden_states.size()[1]
|
||||
|
||||
transfm.update_rotary_pos_emb_cache(kv_seq_len, ntk_alpha=1.0)
|
||||
cos, sin = transfm._rotary_pos_emb_cache
|
||||
rotary_pos_emb_list = [[cos[:, :kv_seq_len], sin[:, :kv_seq_len]]]
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
hidden_states = transfm.drop(hidden_states)
|
||||
output_shape = input_shape + (hidden_states.size(-1),)
|
||||
|
||||
def _norm(self, x):
|
||||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
for block in transfm.h:
|
||||
hidden_states = forwardQWenBlock(block, hidden_states, rotary_pos_emb_list=rotary_pos_emb_list)
|
||||
|
||||
def forward(self, x):
|
||||
output = self._norm(x.float()).type_as(x)
|
||||
return output * self.weight
|
||||
hidden_states = transfm.ln_f(hidden_states)
|
||||
hidden_states = hidden_states.view(output_shape)
|
||||
|
||||
lm_logits = qwen.lm_head(hidden_states)
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
labels = labels.to(lm_logits.device)
|
||||
shift_logits = lm_logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
loss_fct = CrossEntropyLoss()
|
||||
loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
|
||||
|
||||
# shift_labels = torch.ones([1,19]).to(lm_logits.device).to(torch.int64)
|
||||
# shift_logits = lm_logits[..., :-1, :].contiguous()
|
||||
# loss_fct = CrossEntropyLoss()
|
||||
# loss = loss_fct(
|
||||
# shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
|
||||
# )
|
||||
# loss.backward()
|
||||
|
||||
return lm_logits
|
||||
|
|
|
|||
Loading…
Reference in New Issue