Witllm/qwen/modeling_qwen.py

import copy
import math
import inspect
import os
import gc
from tqdm import auto as tqdm_lib
import json
from typing import TYPE_CHECKING, Optional, Tuple, Union, Callable, List, Any, Generator

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 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

# tracker = mem_tracker.MemTracker()
# tracker.track()


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):
        new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
        tensor = tensor.view(new_shape)
        return tensor

    def _merge_heads(self, tensor, num_heads, attn_head_size):
        tensor = tensor.contiguous()
        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):
        super().__init__()
        ff_dim_in = config.intermediate_size // 2
        self.w1 = nn.Linear(config.hidden_size, ff_dim_in, bias=not config.no_bias)
        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,
            eps=config.layer_norm_epsilon,
        )
        self.attn = QWenAttention(config, index)
        self.ln_2 = RMSNorm(
            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):
    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)

        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)

        self.h = nn.ModuleList([QWenBlock(config, i) for i in range(config.num_hidden_layers)])
        self.ln_f = RMSNorm(
            self.embed_dim,
            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)]

        hidden_states = self.drop(hidden_states)
        output_shape = input_shape + (hidden_states.size(-1),)

        all_hidden_states = None
        for block in self.h:
            hidden_states = block(hidden_states, rotary_pos_emb_list=rotary_pos_emb_list)

        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)


class QWenLMHeadModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        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}")
        with open(resolved_archive_file, "r") as f:
            index = json.loads(f.read())
        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):
        metadata = getattr(state_dict, "_metadata", None)
        state_dict = state_dict.copy()
        if metadata is not None:
            state_dict._metadata = metadata
        error_msgs = []

        def load(module: nn.Module, state_dict, prefix=""):
            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
            args = (state_dict, prefix, local_metadata, True, [], [], error_msgs)
            if len([key for key in state_dict if key.startswith(prefix)]) > 0:
                module._load_from_state_dict(*args)

            for name, child in module._modules.items():
                if child is not None:
                    load(child, state_dict, prefix + name + ".")

        load(model_to_load, state_dict, prefix=start_prefix)
        del state_dict
        return error_msgs

    def _load_pretrained_model(cls, resolved_archive_file):
        start_prefix = ""
        model_to_load = cls
        if len(resolved_archive_file) > 1:
            resolved_archive_file = tqdm_lib.tqdm(resolved_archive_file, desc="Loading checkpoint shards")
        for shard_file in resolved_archive_file:
            state_dict = safe_load_file(shard_file)
            cls._load_state_dict_into_model(model_to_load, state_dict, start_prefix)
            del state_dict  # force memory release
            gc.collect()
        print(f"All model checkpoint weights were used when initializing {cls.__class__.__name__}.\n")
        return cls

    @torch.no_grad()
    def chat(
        self,
        tokenizer: PreTrainedTokenizer,
        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,
        )
        decoded, response, end_reason = decode_tokens(
            outputs[0],
            tokenizer,
            raw_text_len=len(raw_text),
            context_length=len(context_tokens),
            errors="replace",
        )
        history.append((query, response))
        return response, history, decoded

    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 = 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, :]

            # repetition_penalty
            penalty = self.config.repetition_penalty
            score = torch.gather(next_token_scores, 1, input_ids)
            # if score < 0 then repetition penalty has to be multiplied to reduce the token probabilities
            score = torch.where(score < 0, score * penalty, score / penalty)
            next_token_scores = next_token_scores.scatter_(1, input_ids, score)

            # top_p
            top_p = self.config.top_p
            filter_value = -float("Inf")
            min_tokens_to_keep = 1
            sorted_logits, sorted_indices = torch.sort(next_token_scores, descending=False)
            cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
            # Remove tokens with cumulative top_p above the threshold (token with 0 are kept)
            sorted_indices_to_remove = cumulative_probs <= (1 - top_p)
            # Keep at least min_tokens_to_keep
            sorted_indices_to_remove[..., -min_tokens_to_keep:] = 0
            # scatter sorted tensors to original indexing
            indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
            next_token_scores = next_token_scores.masked_fill(indices_to_remove, filter_value)

            # sample
            probs = nn.functional.softmax(next_token_scores, dim=-1)
            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)

            next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)

            # update generated ids, model inputs, and length for next step
            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)

            unfinished_sequences = unfinished_sequences.mul(
                next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
            )

            # decoded, response, end_reason = decode_tokens(
            #     next_tokens,
            #     tokenizer,
            #     raw_text_len=0,
            #     context_length=0,
            #     errors="replace",
            # )
            # print(decoded)

            # stop when each sentence is finished
            if unfinished_sequences.max() == 0:
                this_peer_finished = True

            if this_peer_finished:
                break
        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 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)

            emb = torch.cat((freqs, freqs), dim=-1)
            emb = rearrange(emb, "n d -> 1 n 1 d")

            cos, sin = emb.cos(), emb.sin()
            self._rotary_pos_emb_cache = [cos, sin]

    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]]


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 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)


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):
        output = self._norm(x.float()).type_as(x)
        return output * self.weight