diff --git a/qwen/demo.py b/qwen/demo.py index 250fe7b..edcb1f3 100644 --- a/qwen/demo.py +++ b/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 () diff --git a/qwen/modeling_qwen.py b/qwen/modeling_qwen.py index da27cd9..1e85ebd 100644 --- a/qwen/modeling_qwen.py +++ b/qwen/modeling_qwen.py @@ -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