colin
/
Witllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

PreTrainedModel to mm.Module

This commit is contained in:
Colin 2024-01-20 20:04:45 +08:00
parent 0458e7303c
commit 12dcbec718
4 changed files with 103 additions and 92 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
qwen/config.json
View File

@ -33,5 +33,14 @@
"use_dynamic_ntk": true,
"use_flash_attn": "auto",
"use_logn_attn": true,
"vocab_size": 151936
"vocab_size": 151936,
"chat_format": "chatml",
"eos_token_id": 151643,
"pad_token_id": 151643,
"max_window_size": 6144,
"max_new_tokens": 512,
"do_sample": true,
"top_k": 0,
"top_p": 0.8,
"repetition_penalty": 1.1
}

5
qwen/demo.py
View File

@ -52,7 +52,8 @@ print(model)
tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
model = model.from_pretrained(model_dir, config=config, device_map="auto", trust_remote_code=True)
# model = model.from_pretrained(model_dir, config=config, device_map="cuda:1", trust_remote_code=True)
model = model.from_pretrained(model_dir).cuda()
# model = model.eval()
model = model.train() # control by @torch.no_grad()
@ -80,7 +81,7 @@ print(decode_tokens)
# 第二轮对话
response, history, decode_tokens = model.chat(tokenizer, "给我讲一个年轻人奋斗创业最终取得成功的故事。", "", history=None)
print(response)
print(decode_tokens)
# <|im_start|>system

12
qwen/generation_config.json
View File

@ -1,12 +0,0 @@
{
"chat_format": "chatml",
"eos_token_id": 151643,
"pad_token_id": 151643,
"max_window_size": 6144,
"max_new_tokens": 512,
"do_sample": true,
"top_k": 0,
"top_p": 0.8,
"repetition_penalty": 1.1,
"transformers_version": "4.31.0"
}

165
qwen/modeling_qwen.py
View File

@ -6,6 +6,10 @@
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
@ -37,6 +41,11 @@ from qwen_generation_utils import (
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("..")
@ -192,7 +201,7 @@ class QWenBlock(nn.Module):
return hidden_states
class QWenPreTrainedModel(PreTrainedModel):
class QWenPreTrainedModel(nn.Module):
config_class = QWenConfig
base_model_prefix = "transformer"
is_parallelizable = False
@ -200,7 +209,7 @@ class QWenPreTrainedModel(PreTrainedModel):
_no_split_modules = ["QWenBlock"]
def __init__(self, *inputs, **kwargs):
super().__init__(*inputs, **kwargs)
super().__init__()
class QWenModel(QWenPreTrainedModel):
@ -231,8 +240,6 @@ class QWenModel(QWenPreTrainedModel):
eps=config.layer_norm_epsilon,
)
self.post_init()
def get_ntk_alpha(self, true_seq_len):
context_value = math.log(true_seq_len / self.seq_length, 2) + 1
ntk_alpha = 2 ** math.ceil(context_value) - 1
@ -242,7 +249,6 @@ class QWenModel(QWenPreTrainedModel):
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
):
if input_ids is not None and inputs_embeds is not None:
@ -257,8 +263,6 @@ class QWenModel(QWenPreTrainedModel):
else:
raise ValueError("You have to specify either input_ids or inputs_embeds")
head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
if inputs_embeds is None:
inputs_embeds = self.wte(input_ids)
hidden_states = inputs_embeds
@ -288,13 +292,14 @@ class QWenModel(QWenPreTrainedModel):
return BaseModelOutputWithPast(last_hidden_state=hidden_states, hidden_states=all_hidden_states)
class QWenLMHeadModel(QWenPreTrainedModel):
class QWenLMHeadModel(nn.Module):
def __init__(self, config):
super().__init__(config)
super().__init__()
self.config = config
self.transformer = QWenModel(config)
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
self.post_init()
self.generation_config = GenerationConfig.from_model_config(config)
def prepare_inputs_for_generation(self, input_ids, inputs_embeds=None, **kwargs):
model_inputs = {"input_ids": input_ids}
@ -303,13 +308,11 @@ class QWenLMHeadModel(QWenPreTrainedModel):
def forward(
self,
input_ids: Optional[torch.LongTensor] = None,
head_mask: Optional[torch.FloatTensor] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
labels: Optional[torch.LongTensor] = None,
) -> Union[Tuple, CausalLMOutputWithPast]:
transformer_outputs = self.transformer(
input_ids,
head_mask=head_mask,
inputs_embeds=inputs_embeds,
)
hidden_states = transformer_outputs[0]
@ -339,6 +342,59 @@ class QWenLMHeadModel(QWenPreTrainedModel):
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
error_msgs = []
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)
error_msgs += 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,
@ -358,15 +414,10 @@ class QWenLMHeadModel(QWenPreTrainedModel):
stop_words_ids = []
max_window_size = kwargs.get("max_window_size", None)
if max_window_size is None:
max_window_size = generation_config.max_window_size
raw_text, context_tokens = make_context(
tokenizer, query, query_assistant, history=history, system=system, max_window_size=max_window_size
)
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(self.device)
input_ids = torch.tensor([context_tokens]).to(next(self.parameters()).device)
outputs = self.generate(
input_ids,
stop_words_ids=stop_words_ids,
@ -385,7 +436,7 @@ class QWenLMHeadModel(QWenPreTrainedModel):
def generate(
self,
inputs: Optional[torch.Tensor] = None,
input_ids: Optional[torch.Tensor] = None,
stop_words_ids=[],
tokenizer=None,
prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
@ -393,13 +444,9 @@ class QWenLMHeadModel(QWenPreTrainedModel):
) -> Union[GenerateOutput, torch.LongTensor]:
generation_config = self.generation_config
# 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
self._validate_model_class()
generation_config = copy.deepcopy(generation_config)
model_kwargs = generation_config.update(**kwargs) # All unused kwargs must be model kwargs
generation_config.validate()
self._validate_model_kwargs(model_kwargs.copy())
# 2. Set generation parameters if not already defined
@ -409,57 +456,14 @@ class QWenLMHeadModel(QWenPreTrainedModel):
eos_token_id = eos_token_id[0]
generation_config.pad_token_id = eos_token_id
# 3. Define model inputs
inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(
inputs, generation_config.bos_token_id, model_kwargs
)
# 4. Define other model kwargs
# 5. Prepare `input_ids` which will be used for auto-regressive generation
input_ids = inputs_tensor if model_input_name == "input_ids" else model_kwargs.pop("input_ids")
# 6. Prepare `max_length` depending on other stopping criteria.
input_ids_length = input_ids.shape[-1]
has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
generation_config.max_length = generation_config.max_new_tokens + input_ids_length
self._validate_generated_length(generation_config, input_ids_length, has_default_max_length)
stop_words_logits_processor = StopWordsLogitsProcessor(
stop_words_ids=stop_words_ids,
eos_token_id=generation_config.eos_token_id,
)
logits_processor = LogitsProcessorList([stop_words_logits_processor])
logits_processor = self._get_logits_processor(
generation_config=generation_config,
input_ids_seq_length=input_ids_length,
encoder_input_ids=inputs_tensor,
prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
logits_processor=logits_processor,
model_kwargs=model_kwargs,
negative_prompt_ids=None,
negative_prompt_attention_mask=None,
)
# 12. expand input_ids with `num_return_sequences` additional sequences per batch
input_ids, model_kwargs = self._expand_inputs_for_generation(
input_ids=input_ids,
expand_size=generation_config.num_return_sequences,
is_encoder_decoder=False,
**model_kwargs,
)
# 13. run sample
pad_token_id = generation_config.pad_token_id
eos_token_id_tensor = torch.tensor([generation_config.eos_token_id]).to(input_ids.device)
# init values
stopping_criteria = self._get_stopping_criteria(
generation_config=generation_config, stopping_criteria=StoppingCriteriaList()
)
logits_warper = self._get_logits_warper(generation_config)
# init attention / hidden states / scores tuples
scores = None
@ -469,7 +473,6 @@ class QWenLMHeadModel(QWenPreTrainedModel):
this_peer_finished = False
# auto-regressive generation
while True:
# prepare model inputs
model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
# forward pass to get next token
@ -477,9 +480,24 @@ class QWenLMHeadModel(QWenPreTrainedModel):
next_token_scores = outputs.logits[:, -1, :]
# pre-process distribution
next_token_scores = logits_processor(input_ids, next_token_scores)
next_token_scores = logits_warper(input_ids, next_token_scores)
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 = 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)
@ -489,7 +507,6 @@ class QWenLMHeadModel(QWenPreTrainedModel):
# update generated ids, model inputs, and length for next step
input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
model_kwargs = self._update_model_kwargs_for_generation(outputs, model_kwargs, is_encoder_decoder=False)
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)
@ -508,10 +525,6 @@ class QWenLMHeadModel(QWenPreTrainedModel):
if unfinished_sequences.max() == 0:
this_peer_finished = True
# stop if we exceed the maximum length
if stopping_criteria(input_ids, scores):
this_peer_finished = True
if this_peer_finished:
break
return input_ids