qwen/modeling_qwen.py

@@ -200,7 +200,10 @@ class QwenRunner:
         history = copy.deepcopy(history)
         raw_text, context_tokens = self.prepareInput(tokenizer, query, query_assistant, history, system)
         input_ids = torch.tensor([context_tokens]).to(next(qwen.parameters()).device)
-        self.unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
+        eos_token_id_tensor = torch.tensor([qwen.config.eos_token_id]).to(input_ids.device)
+        pad_token_id = qwen.config.pad_token_id
+
+        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
         while True:
             outputs = self.forwardQWen(input_ids)
             next_token_scores = outputs[:, -1, :]
@@ -208,10 +211,14 @@ class QwenRunner:
             next_token_scores = self.repetition_penalty(input_ids, next_token_scores)
             next_token_scores = self.top_p(next_token_scores)
             next_tokens = self.sample(next_token_scores)
-            finish, next_tokens = self.isFinish(next_tokens)
-            if finish:
+
+            next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+            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)
+            )
+            if unfinished_sequences.max() == 0:
                 break
-            input_ids = torch.cat([input_ids, next_tokens], dim=-1)
 
         decoded, response, end_reason = decode_tokens(
             input_ids[0],
@@ -377,13 +384,3 @@ class QwenRunner:
         probs = nn.functional.softmax(next_token_scores, dim=-1)
         next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
         return next_tokens
-
-    def isFinish(self, next_tokens):
-        pad_token_id = self.qwen.config.pad_token_id
-        eos_token_id_tensor = torch.tensor([self.qwen.config.eos_token_id]).to(next_tokens.device)
-
-        next_tokens = next_tokens * self.unfinished_sequences + pad_token_id * (1 - self.unfinished_sequences)
-        self.unfinished_sequences = self.unfinished_sequences.mul(
-            next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
-        )
-        return self.unfinished_sequences.max() == 0, next_tokens[:, None]

qwen/research_attention.py

@@ -40,14 +40,12 @@ print(model)
 tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
 
 model = model.from_pretrained(model_dir)
-if torch.cuda.device_count() > 0:
-    model = model.cuda()
 model = model.eval()
 
 
 def Dump_tokens_list(model):
     tokens = []
-    for token in range(151851):
+    for token in range(config.eos_token_id):
         decoded, response, end_reason = decode_tokens(
             [token],
             tokenizer,
@@ -72,14 +70,11 @@ def Dump_lm_head_weight(model):
 
 # Dump_lm_head_weight(model)
 
-qk_seq = []
-qk_index = None
+qk_sum = []
+qk_index = []
 
 
 def DumpQK(query, key, causal_mask, index):
-    global qk_seq
-    global qk_index
-    size = query.shape[2]
     scale_factor = 1 / math.sqrt(query.size(-1))
     attn_weight = query @ key.transpose(-2, -1) * scale_factor
     attn_mask = torch.ones(causal_mask.shape, dtype=query.dtype, device=query.device)
@@ -87,11 +82,12 @@ def DumpQK(query, key, causal_mask, index):
     attn_weight = attn_weight * attn_mask
     attn_weight = torch.softmax(attn_weight, dim=-1)
     attn_weight = attn_weight * attn_mask
+    size = query.shape[2]
     qk = attn_weight[0]
     # prePath = "./temp/" + "q@k_seq_" + str(size) + "_layer_" + str(index) + ".png"
     # show.DumpTensorToImage(qk, prePath, GridValue=255)
-    qk_seq.append(qk)
-    qk_index = size
+    qk_sum.append(qk.sum(0))
+    qk_index.append(size)
 
 
 class ResearchRunner(QwenRunner):
@@ -110,6 +106,14 @@ class ResearchRunner(QwenRunner):
         attn_output = attention.c_proj(context_layer)
         return attn_output
 
+    def sample(self, next_token_scores):
+        qk_sum_cat = torch.stack(qk_sum, 0)
+        qk_sum.clear()
+        prePath = "./temp/" + "q@k_sum_seq_" + str(qk_index[-1]) + ".png"
+        show.DumpTensorToImage(qk_sum_cat, prePath, GridValue=255)
+
+        return super().sample(next_token_scores)
+
     def prepareInput(self, tokenizer, query, query_assistant, history, system):
         start_to = [151644]
         n_to = [198]
@@ -124,21 +128,10 @@ class ResearchRunner(QwenRunner):
 
         tokens = system_token + user_token + aassistant_token
         tokens = user_token + aassistant_token
-        tokens = start_to + tokenizer.encode("user\nHi你好\nassistant\n", allowed_special=set())
+        tokens = start_to + tokenizer.encode("user\n你好\nassistant\n", allowed_special=set())
 
         return "", tokens
 
-    def isFinish(self, next_tokens):
-        global qk_seq
-        finish, next = super().isFinish(next_tokens)
-        if finish:
-            for i, s in enumerate(qk_seq):
-                prePath = "./temp/" + "q@k_layer_" + str(i) + ".png"
-                show.DumpTensorToImage(s, prePath, GridValue=255)
-        else:
-            qk_seq = []
-        return finish, next
-
 
 runner = ResearchRunner(model)
 

qwen/research_token.py

@@ -1,136 +0,0 @@
-import torch
-import sys
-
-# from modelscope import snapshot_download
-from transformers import AutoTokenizer
-from transformers import AutoConfig
-
-from modeling_qwen import QWenLMHeadModel
-from modeling_qwen import QwenRunner
-
-
-from qwen_generation_utils import (
-    make_context,
-    decode_tokens,
-)
-
-import torch.nn.functional as F
-
-sys.path.append("..")
-from tools import show
-
-seed = 4321
-torch.manual_seed(seed)
-torch.cuda.manual_seed_all(seed)
-
-# model_dir = snapshot_download("qwen/Qwen-1_8B-Chat")
-model_dir = "/home/colin/.cache/modelscope/hub/qwen/Qwen-1_8B-Chat"
-
-config, kwargs = AutoConfig.from_pretrained(
-    "./",
-    return_unused_kwargs=True,
-    trust_remote_code=True,
-    code_revision=None,
-    _commit_hash=None,
-)
-
-model = QWenLMHeadModel(config)
-print(model)
-
-tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
-
-model = model.from_pretrained(model_dir)
-if torch.cuda.device_count() > 0:
-    model = model.cuda()
-model = model.eval()
-
-
-class ResearchRunner(QwenRunner):
-    def __init__(self, model):
-        super().__init__(model)
-
-    def forwardQWen(
-        self,
-        input_ids=None,
-        labels=None,
-    ):
-        transfm = self.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]]]
-
-        hidden_states = transfm.drop(hidden_states)
-        output_shape = input_shape + (hidden_states.size(-1),)
-
-        for block in transfm.h:
-            self.forwardQWenBlock(block, hidden_states, rotary_pos_emb_list=rotary_pos_emb_list)
-            break
-
-    def forwardQWenBlock(
-        self,
-        block,
-        hidden_states,
-        rotary_pos_emb_list=None,
-    ):
-        layernorm_output = block.ln_1(hidden_states)
-        self.forwardAttention(block.attn, layernorm_output, rotary_pos_emb_list)
-
-    def attention(self, attention, query, key, value, causal_mask):
-        query = query.permute(0, 2, 1, 3)
-        key = key.permute(0, 2, 1, 3)
-        value = value.permute(0, 2, 1, 3)
-        global q
-        global k
-        query = query[:, head_group_index, :, :]
-        key = key[:, head_group_index, :, :]
-        q = torch.cat([q, query], 1)
-        k = torch.cat([k, key], 1)
-
-
-head_group_index = 0
-total_token = 151851
-topk = 10
-
-tokens_str = []
-for token in range(total_token):
-    decoded, response, end_reason = decode_tokens(
-        [token],
-        tokenizer,
-        raw_text_len=0,
-        context_length=0,
-        errors="replace",
-    )
-    tokens_str.append(repr(decoded))
-
-patch_end = list(range(0, total_token, 1000))
-patch_end = patch_end[1:] + [total_token]
-patch_start = 0
-q = torch.zeros((1, 0, 128), dtype=float).to(next(model.parameters()).device)
-k = torch.zeros((1, 0, 128), dtype=float).to(next(model.parameters()).device)
-for end in patch_end:
-    tokens = list(range(patch_start, end))
-    patch_start = end
-    input_ids = torch.tensor([tokens]).to(next(model.parameters()).device)
-    runner = ResearchRunner(model)
-    runner.forwardQWen(input_ids)
-
-q = q[0, :, :]
-k = k[0, :, :].permute(1, 0)
-token_topk = []
-for i in range(total_token):
-    subq = q[i, :]
-    qk = subq @ k
-    values, indices = torch.topk(qk, topk)
-    item = str(i).zfill(7) + " " + tokens_str[i] + " : "
-    for index in indices:
-        item += tokens_str[index] + "  "
-    token_topk.append(item)
-
-show.DumpListToFile(token_topk, "./temp/qwen_token_qk_topk_head_group_" + str(head_group_index) + ".txt")
-
-print("decoded")