Compare commits
2 Commits
3f296ccdb2
...
b7c27af6c8
Author | SHA1 | Date |
---|---|---|
Colin | b7c27af6c8 | |
Colin | 185278f3a9 |
|
@ -200,10 +200,7 @@ 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)
|
||||
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)
|
||||
self.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, :]
|
||||
|
@ -211,14 +208,10 @@ 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)
|
||||
|
||||
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:
|
||||
finish, next_tokens = self.isFinish(next_tokens)
|
||||
if finish:
|
||||
break
|
||||
input_ids = torch.cat([input_ids, next_tokens], dim=-1)
|
||||
|
||||
decoded, response, end_reason = decode_tokens(
|
||||
input_ids[0],
|
||||
|
@ -384,3 +377,13 @@ 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]
|
||||
|
|
|
@ -40,12 +40,14 @@ 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(config.eos_token_id):
|
||||
for token in range(151851):
|
||||
decoded, response, end_reason = decode_tokens(
|
||||
[token],
|
||||
tokenizer,
|
||||
|
@ -70,11 +72,14 @@ def Dump_lm_head_weight(model):
|
|||
|
||||
# Dump_lm_head_weight(model)
|
||||
|
||||
qk_sum = []
|
||||
qk_index = []
|
||||
qk_seq = []
|
||||
qk_index = None
|
||||
|
||||
|
||||
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)
|
||||
|
@ -82,12 +87,11 @@ 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_sum.append(qk.sum(0))
|
||||
qk_index.append(size)
|
||||
qk_seq.append(qk)
|
||||
qk_index = size
|
||||
|
||||
|
||||
class ResearchRunner(QwenRunner):
|
||||
|
@ -106,14 +110,6 @@ 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]
|
||||
|
@ -128,10 +124,21 @@ class ResearchRunner(QwenRunner):
|
|||
|
||||
tokens = system_token + user_token + aassistant_token
|
||||
tokens = user_token + aassistant_token
|
||||
tokens = start_to + tokenizer.encode("user\n你好\nassistant\n", allowed_special=set())
|
||||
tokens = start_to + tokenizer.encode("user\nHi你好\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)
|
||||
|
||||
|
|
|
@ -0,0 +1,136 @@
|
|||
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")
|
Loading…
Reference in New Issue