learn2pro commited on
Commit
098b164
·
verified ·
1 Parent(s): 65f5fb4

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
added_tokens.json ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "</think>": 151668,
3
+ "</tool_call>": 151658,
4
+ "</tool_response>": 151666,
5
+ "<think>": 151667,
6
+ "<tool_call>": 151657,
7
+ "<tool_response>": 151665,
8
+ "<|box_end|>": 151649,
9
+ "<|box_start|>": 151648,
10
+ "<|endoftext|>": 151643,
11
+ "<|file_sep|>": 151664,
12
+ "<|fim_middle|>": 151660,
13
+ "<|fim_pad|>": 151662,
14
+ "<|fim_prefix|>": 151659,
15
+ "<|fim_suffix|>": 151661,
16
+ "<|im_end|>": 151645,
17
+ "<|im_start|>": 151644,
18
+ "<|image_pad|>": 151655,
19
+ "<|object_ref_end|>": 151647,
20
+ "<|object_ref_start|>": 151646,
21
+ "<|quad_end|>": 151651,
22
+ "<|quad_start|>": 151650,
23
+ "<|repo_name|>": 151663,
24
+ "<|video_pad|>": 151656,
25
+ "<|vision_end|>": 151653,
26
+ "<|vision_pad|>": 151654,
27
+ "<|vision_start|>": 151652
28
+ }
chat_template.jinja ADDED
@@ -0,0 +1,89 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {%- if tools %}
2
+ {{- '<|im_start|>system\n' }}
3
+ {%- if messages[0].role == 'system' %}
4
+ {{- messages[0].content + '\n\n' }}
5
+ {%- endif %}
6
+ {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
7
+ {%- for tool in tools %}
8
+ {{- "\n" }}
9
+ {{- tool | tojson }}
10
+ {%- endfor %}
11
+ {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
12
+ {%- else %}
13
+ {%- if messages[0].role == 'system' %}
14
+ {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
15
+ {%- endif %}
16
+ {%- endif %}
17
+ {%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
18
+ {%- for message in messages[::-1] %}
19
+ {%- set index = (messages|length - 1) - loop.index0 %}
20
+ {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
21
+ {%- set ns.multi_step_tool = false %}
22
+ {%- set ns.last_query_index = index %}
23
+ {%- endif %}
24
+ {%- endfor %}
25
+ {%- for message in messages %}
26
+ {%- if message.content is string %}
27
+ {%- set content = message.content %}
28
+ {%- else %}
29
+ {%- set content = '' %}
30
+ {%- endif %}
31
+ {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
32
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
33
+ {%- elif message.role == "assistant" %}
34
+ {%- set reasoning_content = '' %}
35
+ {%- if message.reasoning_content is string %}
36
+ {%- set reasoning_content = message.reasoning_content %}
37
+ {%- else %}
38
+ {%- if '</think>' in content %}
39
+ {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
40
+ {%- set content = content.split('</think>')[-1].lstrip('\n') %}
41
+ {%- endif %}
42
+ {%- endif %}
43
+ {%- if loop.index0 > ns.last_query_index %}
44
+ {%- if loop.last or (not loop.last and reasoning_content) %}
45
+ {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
46
+ {%- else %}
47
+ {{- '<|im_start|>' + message.role + '\n' + content }}
48
+ {%- endif %}
49
+ {%- else %}
50
+ {{- '<|im_start|>' + message.role + '\n' + content }}
51
+ {%- endif %}
52
+ {%- if message.tool_calls %}
53
+ {%- for tool_call in message.tool_calls %}
54
+ {%- if (loop.first and content) or (not loop.first) %}
55
+ {{- '\n' }}
56
+ {%- endif %}
57
+ {%- if tool_call.function %}
58
+ {%- set tool_call = tool_call.function %}
59
+ {%- endif %}
60
+ {{- '<tool_call>\n{"name": "' }}
61
+ {{- tool_call.name }}
62
+ {{- '", "arguments": ' }}
63
+ {%- if tool_call.arguments is string %}
64
+ {{- tool_call.arguments }}
65
+ {%- else %}
66
+ {{- tool_call.arguments | tojson }}
67
+ {%- endif %}
68
+ {{- '}\n</tool_call>' }}
69
+ {%- endfor %}
70
+ {%- endif %}
71
+ {{- '<|im_end|>\n' }}
72
+ {%- elif message.role == "tool" %}
73
+ {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
74
+ {{- '<|im_start|>user' }}
75
+ {%- endif %}
76
+ {{- '\n<tool_response>\n' }}
77
+ {{- content }}
78
+ {{- '\n</tool_response>' }}
79
+ {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
80
+ {{- '<|im_end|>\n' }}
81
+ {%- endif %}
82
+ {%- endif %}
83
+ {%- endfor %}
84
+ {%- if add_generation_prompt %}
85
+ {{- '<|im_start|>assistant\n' }}
86
+ {%- if enable_thinking is defined and enable_thinking is false %}
87
+ {{- '<think>\n\n</think>\n\n' }}
88
+ {%- endif %}
89
+ {%- endif %}
config.json ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "BuddyGPTForCausalLM"
4
+ ],
5
+ "attention_dropout": 0.0,
6
+ "eos_token_id": 151645,
7
+ "hidden_act": "silu",
8
+ "hidden_size": 768,
9
+ "initializer_range": 0.02,
10
+ "intermediate_size": 1536,
11
+ "model_type": "buddygpt",
12
+ "num_attention_heads": 16,
13
+ "num_hidden_layers": 8,
14
+ "num_key_value_heads": 8,
15
+ "num_seq_len": 1024,
16
+ "pad_token_id": 151643,
17
+ "rms_norm_eps": 1e-06,
18
+ "rope_theta": 10000.0,
19
+ "tie_word_embeddings": true,
20
+ "torch_dtype": "float32",
21
+ "transformers_version": "4.52.4",
22
+ "use_cache": true,
23
+ "vocab_size": 151669
24
+ }
configuration_buddygpt.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers.configuration_utils import PretrainedConfig
2
+ from loguru import logger as logging
3
+
4
+
5
+ class BuddyGPTConfig(PretrainedConfig):
6
+ """ TinyLLM 配置文件
7
+ """
8
+
9
+ model_type = "buddygpt"
10
+ keys_to_ignore_at_inference = ["past_key_values"]
11
+
12
+ def __init__(
13
+ self,
14
+ vocab_size=151669,
15
+ hidden_size=4096,
16
+ intermediate_size=4096,
17
+ num_hidden_layers=32,
18
+ num_attention_heads=32,
19
+ num_key_value_heads=None,
20
+ hidden_act="silu",
21
+ num_seq_len=2048,
22
+ initializer_range=0.02,
23
+ rms_norm_eps=1e-6,
24
+ use_cache=True,
25
+ pad_token_id=None,
26
+ bos_token_id=None,
27
+ eos_token_id=None,
28
+ tie_word_embeddings=False,
29
+ rope_theta=10000.0,
30
+ attention_dropout=0.0,
31
+ _attn_implementation="sdpa",
32
+ **kwargs
33
+ ):
34
+ self.vocab_size = vocab_size
35
+ self.num_seq_len = num_seq_len
36
+ self.hidden_size = hidden_size
37
+ self.intermediate_size = intermediate_size
38
+ self.num_hidden_layers = num_hidden_layers
39
+ self.num_attention_heads = num_attention_heads
40
+
41
+ # for backward compatibility
42
+ if num_key_value_heads is None:
43
+ num_key_value_heads = num_attention_heads
44
+
45
+ self.num_key_value_heads = num_key_value_heads
46
+ self.hidden_act = hidden_act
47
+ self.initializer_range = initializer_range
48
+ self.rms_norm_eps = rms_norm_eps
49
+ self.use_cache = use_cache
50
+ self.rope_theta = rope_theta
51
+ self.attention_dropout = attention_dropout
52
+ self._attn_implementation = _attn_implementation
53
+
54
+ super().__init__(
55
+ pad_token_id=pad_token_id,
56
+ bos_token_id=bos_token_id,
57
+ eos_token_id=eos_token_id,
58
+ tie_word_embeddings=tie_word_embeddings,
59
+ **kwargs
60
+ )
generation_config.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "eos_token_id": 151645,
4
+ "pad_token_id": 151643,
5
+ "transformers_version": "4.52.4"
6
+ }
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1d8150dd7b3a3c451da97bfd776400aaed387dcec7b9884c6a5529dd64d44478
3
+ size 635908536
modeling_buddygpt.py ADDED
@@ -0,0 +1,966 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Tiny LLM 模型架构
3
+
4
+ 到处抄,整体还是Llama2的模型架构
5
+ """
6
+
7
+ import math
8
+ import warnings
9
+ from threading import Thread
10
+ from typing import List, Optional, Tuple, Union
11
+
12
+ import torch
13
+ import torch.nn.functional as F
14
+ import torch.utils.checkpoint
15
+ from torch import nn
16
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
17
+
18
+ from transformers.activations import ACT2FN
19
+ from transformers.cache_utils import Cache, DynamicCache
20
+ from transformers.modeling_attn_mask_utils import (
21
+ _prepare_4d_causal_attention_mask,
22
+ _prepare_4d_causal_attention_mask_for_sdpa,
23
+ )
24
+ from transformers.modeling_outputs import (
25
+ BaseModelOutputWithPast,
26
+ CausalLMOutputWithPast,
27
+ SequenceClassifierOutputWithPast,
28
+ )
29
+ from transformers.modeling_utils import PreTrainedModel
30
+ from transformers.generation.utils import GenerationConfig, GenerationMixin
31
+ from transformers.generation.logits_process import LogitsProcessorList
32
+
33
+ from model.configuration_buddygpt import BuddyGPTConfig
34
+ from loguru import logger
35
+
36
+ from model.generation_utils import TextIterStreamer, make_context, OutputRepetitionPenaltyLogitsProcessor, parse_pot_no_stream
37
+
38
+
39
+ def report_memory(name):
40
+ """Simple GPU memory report."""
41
+ mega_bytes = 1024.0 * 1024.0
42
+ string = name + " memory (MB)"
43
+ # 变量分配显存
44
+ string += " | allocated: {}".format(torch.cuda.memory_allocated() / mega_bytes)
45
+ string += " | max allocated: {}".format(
46
+ torch.cuda.max_memory_allocated() / mega_bytes
47
+ )
48
+ # 缓存和变量分配显存,实际显存还需要+pytorch context
49
+ string += " | reserved: {}".format(torch.cuda.memory_reserved() / mega_bytes)
50
+ string += " | max reserved: {}".format(
51
+ torch.cuda.max_memory_reserved() / mega_bytes
52
+ )
53
+ try:
54
+ if torch.distributed.get_rank() == 0:
55
+ print(
56
+ "[Rank {}] {}".format(torch.distributed.get_rank(), string), flush=True
57
+ )
58
+ pass
59
+ except:
60
+ pass
61
+
62
+
63
+ import torch
64
+ import torch.nn as nn
65
+
66
+ class RotaryEmbedding(nn.Module):
67
+ def __init__(self, dim, max_position_embeddings=2048, base=100000, device=None):
68
+ """ 旋转位置编码
69
+ - dim (int): 旋转嵌入的维度大小。
70
+ - max_position_embeddings (int): 预计算的最大位置嵌入数,默认为2048。
71
+ - base (int): 用于计算逆频率的基本频率,默认为10000。
72
+ """
73
+ super().__init__()
74
+
75
+ self.dim = dim
76
+ self.max_position_embeddings = max_position_embeddings
77
+ self.base = base
78
+ # 计算逆频率值,并将其注册为模型的缓冲区
79
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
80
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
81
+
82
+ # 为了支持`torch.jit.trace`功能,立即计算预存储的余弦和正弦缓存
83
+ self._set_cos_sin_cache(seq_len=max_position_embeddings, device=self.inv_freq.device)
84
+
85
+ def _set_cos_sin_cache(self, seq_len, device):
86
+ """ 预计算的余弦和正弦缓存
87
+ """
88
+ self.max_seq_len_cached = seq_len
89
+ # 创建一个从0到最大序列长度-1的整数张量,与 inv_freq 具有相同的设备和数据类型
90
+ t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
91
+
92
+ # 计算每个位置与每个维度的频率,形成频谱矩阵
93
+ # freqs = torch.outer(t, self.inv_freq)
94
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
95
+
96
+ # 不同于论文中的实现,这里采用了不同的排列方式以获得相同的计算结果
97
+ emb = torch.cat((freqs, freqs), dim=-1)
98
+ self.register_buffer("cos_cached", emb.cos(), persistent=False)
99
+ self.register_buffer("sin_cached", emb.sin(), persistent=False)
100
+
101
+ def forward(self, x, seq_len):
102
+ # x: [bs, num_attention_heads, seq_len, head_size]
103
+ if seq_len > self.max_seq_len_cached:
104
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device)
105
+
106
+ return (
107
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
108
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
109
+ )
110
+
111
+
112
+
113
+ # Copied from transformers.models.mistral.modeling_mistral.apply_rotary_pos_emb
114
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
115
+ """ 在 qk 应用旋转位置编码
116
+
117
+ Args:
118
+ q (`torch.Tensor`): q
119
+ k (`torch.Tensor`): k
120
+ cos (`torch.Tensor`): 旋转位置嵌入的余弦部分
121
+ sin (`torch.Tensor`): 旋转位置嵌入的正弦部分
122
+ position_ids (`torch.Tensor`): 与q和k对应位置的标记索引。例如,在处理KV缓存时,可以使用偏移过的位置ID。
123
+ unsqueeze_dim (`int`, *optional*, defaults to 1): 'unsqueeze_dim' 参数指定了沿哪个维度对 cos[position_ids]
124
+ 和 sin[position_ids] 进行扩展,以便它们能够适当地广播到 q 和 k 的���度上。
125
+ 例如,注意 cos[position_ids] 和 sin[position_ids] 具有形状 [batch_size, seq_len, head_dim]。
126
+ 那么,如果 q 和 k 的形状分别为 [batch_size, heads, seq_len, head_dim],
127
+ 则设置 unsqueeze_dim=1 可使 cos[position_ids] 和 sin[position_ids] 可以广播到 q 和 k 的形状上。
128
+ 同样地,如果 q 和 k 的形状为 [batch_size, seq_len, heads, head_dim],则应将 unsqueeze_dim 设置为 2
129
+ Returns:
130
+ 包含使用旋转位置嵌入变换后的q和k张量的 `tuple(torch.Tensor)`。
131
+ """
132
+ def rotate_half(x):
133
+ """ 旋转输入一半的 hidden dim
134
+ """
135
+ x1, x2 = x.chunk(2, dim=-1)
136
+ return torch.cat((-x2, x1), dim=-1)
137
+
138
+ # print("ori cos: ", cos.shape)
139
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
140
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
141
+
142
+ # print("q: ", q.shape)
143
+ # print("cos: ", cos.shape)
144
+ # print("sin: ", sin.shape)
145
+ # print("rotate_half: ", rotate_half(q).shape)
146
+ # e^(i*theta) = cos(theta) + i*sin(theta)
147
+ # 二维变量 x1+i*x2 旋转 theta 角度 (x1+i*x2)*(cos(theta)+i*sin(theta))
148
+ # 旋转后 x1' = x1*cos(theta) - x2*sin(theta)
149
+ # 旋转后 x2' = x1*sin(theta) + x2*cos(theta)
150
+ q_embed = (q * cos) + (rotate_half(q) * sin)
151
+ k_embed = (k * cos) + (rotate_half(k) * sin)
152
+ return q_embed, k_embed
153
+
154
+ class GateMLP(nn.Module):
155
+ def __init__(self, config):
156
+ super().__init__()
157
+ self.config = config
158
+ self.hidden_size = config.hidden_size
159
+ self.intermediate_size = config.intermediate_size
160
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
161
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
162
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
163
+ self.act_fn = nn.SiLU()
164
+
165
+ def forward(self, x):
166
+ intermediate = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
167
+ down_proj = self.down_proj(intermediate)
168
+ return down_proj
169
+
170
+
171
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
172
+ """
173
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
174
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
175
+ """
176
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
177
+ if n_rep == 1:
178
+ return hidden_states
179
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
180
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
181
+
182
+
183
+ class SelfAttention(nn.Module):
184
+ """多头注意力"""
185
+
186
+ def __init__(self, config: BuddyGPTConfig, layer_idx: Optional[int] = None):
187
+ super().__init__()
188
+ self.config = config
189
+ self.layer_idx = layer_idx
190
+ if layer_idx is None:
191
+ logger.warning(
192
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
193
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
194
+ "when creating this class."
195
+ )
196
+
197
+ self.hidden_size = config.hidden_size
198
+ self.num_heads = config.num_attention_heads
199
+ self.head_dim = self.hidden_size // self.num_heads
200
+ self.num_key_value_heads = config.num_key_value_heads
201
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
202
+ self.num_seq_len = config.num_seq_len
203
+ self.rope_theta = config.rope_theta
204
+ # 因果自回归模式
205
+ self.is_causal = True
206
+ self.attention_dropout = config.attention_dropout
207
+
208
+ if (self.head_dim * self.num_heads) != self.hidden_size:
209
+ raise ValueError(
210
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
211
+ f" and `num_heads`: {self.num_heads})."
212
+ )
213
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)
214
+ self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
215
+ self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)
216
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
217
+
218
+ self.rotary_emb = RotaryEmbedding(
219
+ self.head_dim,
220
+ max_position_embeddings=self.num_seq_len,
221
+ base=self.rope_theta,
222
+ )
223
+
224
+ def forward(
225
+ self,
226
+ hidden_states: torch.Tensor,
227
+ attention_mask: Optional[torch.Tensor] = None,
228
+ position_ids: Optional[torch.LongTensor] = None,
229
+ past_key_value: Optional[Cache] = None,
230
+ output_attentions: bool = False,
231
+ use_cache: bool = False,
232
+ **kwargs,
233
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
234
+ if "padding_mask" in kwargs:
235
+ warnings.warn(
236
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
237
+ )
238
+ bsz, q_len, _ = hidden_states.size()
239
+
240
+ query_states = self.q_proj(hidden_states)
241
+ key_states = self.k_proj(hidden_states)
242
+ value_states = self.v_proj(hidden_states)
243
+
244
+ # 重新投影,变成多头注意力结构
245
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
246
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
247
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
248
+
249
+ kv_seq_len = key_states.shape[-2]
250
+ if past_key_value is not None:
251
+ if self.layer_idx is None:
252
+ raise ValueError(
253
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
254
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
255
+ "with a layer index."
256
+ )
257
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
258
+ # 应用旋转位置编码到 qk 向量
259
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
260
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
261
+ # query_states, key_states = self.rotary_emb.apply_rotary_emb(query_states, key_states)
262
+
263
+ # 如果存在缓存,则更新 kv
264
+ if past_key_value is not None:
265
+ key_states, value_states = past_key_value.update(
266
+ key_states, value_states, self.layer_idx
267
+ )
268
+
269
+ # repeat k/v heads if n_kv_heads < n_heads
270
+ # 如果 num_key_value_heads 小于 num_heads,则重复key和value向量以匹配头数量
271
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
272
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
273
+
274
+ # 计算注意力权重
275
+ attn_weights = torch.matmul(
276
+ query_states, key_states.transpose(2, 3)
277
+ ) / math.sqrt(self.head_dim)
278
+
279
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
280
+ raise ValueError(
281
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
282
+ f" {attn_weights.size()}"
283
+ )
284
+
285
+ if attention_mask is not None:
286
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
287
+ raise ValueError(
288
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
289
+ )
290
+
291
+ attn_weights = attn_weights + attention_mask
292
+
293
+ # softmax归一化注意力权重,并转换至float32类型以防止数值溢出
294
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
295
+ attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
296
+ # 注意力输出
297
+ attn_output = torch.matmul(attn_weights, value_states)
298
+
299
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
300
+ raise ValueError(
301
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
302
+ f" {attn_output.size()}"
303
+ )
304
+
305
+ # 还原注意力输出的形状以与后续层对接
306
+ attn_output = attn_output.transpose(1, 2).contiguous()
307
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
308
+
309
+ # 通过o_proj层进一步处理注意力输出
310
+ attn_output = self.o_proj(attn_output)
311
+
312
+ if not output_attentions:
313
+ attn_weights = None
314
+
315
+ return attn_output, attn_weights, past_key_value
316
+
317
+
318
+ class SdpaAttention(SelfAttention):
319
+ """使用 torch.nn.functional.scaled_dot_product_attention 实现的注意力模块。
320
+ 该模块继承自 `SelfAttention`,因为模块的权重保持不变。唯一的变化在于前向传播过程中适应 SDPA API。
321
+ Scaled Dot Product Attention (SDPA)
322
+ """
323
+
324
+ def forward(
325
+ self,
326
+ hidden_states: torch.Tensor,
327
+ attention_mask: Optional[torch.Tensor] = None,
328
+ position_ids: Optional[torch.LongTensor] = None,
329
+ past_key_value: Optional[Cache] = None,
330
+ output_attentions: bool = False,
331
+ use_cache: bool = False,
332
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
333
+ # 当设置output_attentions=True时,由于torch.nn.functional.scaled_dot_product_attention不支持直接返回注意力权重
334
+ # 因此暂时降级回用父类的手动实现方式,并发出警告提示用户未来版本的更改要求
335
+ if output_attentions:
336
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
337
+ logger.warning_once(
338
+ "Model is using SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
339
+ 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
340
+ )
341
+ return super().forward(
342
+ hidden_states=hidden_states,
343
+ attention_mask=attention_mask,
344
+ position_ids=position_ids,
345
+ past_key_value=past_key_value,
346
+ output_attentions=output_attentions,
347
+ use_cache=use_cache,
348
+ )
349
+ # 获取输入维度信息
350
+ bsz, q_len, _ = hidden_states.size() # (bsz, q_len, hidden_dim)
351
+
352
+ # 对输入进行线性映射得到query、key、value向量
353
+ query_states = self.q_proj(hidden_states) # (bsz, q_len, n_head * head_dim)
354
+ key_states = self.k_proj(hidden_states) # (bsz, q_len, n_kv_head * head_dim)
355
+ value_states = self.v_proj(hidden_states) # (bsz, q_len, n_kv_head * head_dim)
356
+
357
+ # 将映射后的向量调整为多头注意力所需格式
358
+ query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) # (bsz, n_head, q_len, head_dim)
359
+ key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) # (bsz, n_kv_head, q_len, head_dim)
360
+ value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) # (bsz, n_kv_head, q_len, head_dim)
361
+
362
+ # 计算有效的 kv 序列长度(考虑缓存的情况)
363
+ kv_seq_len = key_states.shape[-2] # q_len
364
+ if past_key_value is not None:
365
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
366
+
367
+ # 应用旋转位置编码到 qk 向量
368
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
369
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
370
+
371
+
372
+ # 如果有缓存,更新key和value状态
373
+ if past_key_value is not None:
374
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
375
+ key_states, value_states = past_key_value.update(
376
+ key_states, value_states, self.layer_idx, cache_kwargs
377
+ )
378
+
379
+ key_states = repeat_kv(key_states, self.num_key_value_groups) # (bsz, n_head, q_len, head_dim)
380
+ value_states = repeat_kv(value_states, self.num_key_value_groups) # (bsz, n_head, q_len, head_dim)
381
+
382
+ if attention_mask is not None:
383
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
384
+ raise ValueError(
385
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
386
+ )
387
+
388
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
389
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
390
+ if query_states.device.type == "cuda" and attention_mask is not None:
391
+ query_states = query_states.contiguous()
392
+ key_states = key_states.contiguous()
393
+ value_states = value_states.contiguous()
394
+
395
+ # 使用scaled_dot_product_attention进行计算
396
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
397
+ query_states,
398
+ key_states,
399
+ value_states,
400
+ attn_mask=attention_mask,
401
+ dropout_p=self.attention_dropout if self.training else 0.0,
402
+ # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
403
+ is_causal=self.is_causal and attention_mask is None and q_len > 1,
404
+ )
405
+
406
+ # 还原注意力输出的形状
407
+ attn_output = attn_output.transpose(1, 2).contiguous() # (bsz, q_len, n_head, head_dim)
408
+ attn_output = attn_output.view(bsz, q_len, self.hidden_size) # (bsz, q_len, hidden_dim)
409
+
410
+ # 将注意力输出通过最终的线性层(o_proj层)
411
+ attn_output = self.o_proj(attn_output) # (bsz, q_len, hidden_dim)
412
+
413
+ return attn_output, None, past_key_value
414
+
415
+
416
+ class DecoderLayer(nn.Module):
417
+ def __init__(self, config: BuddyGPTConfig, layer_idx: int):
418
+ super().__init__()
419
+ self.hidden_size = config.hidden_size
420
+
421
+ self.self_attn = (SdpaAttention(config, layer_idx) if config._attn_implementation == "sdpa" else SelfAttention(config, layer_idx))
422
+ self.mlp = GateMLP(config)
423
+ self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
424
+ self.post_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
425
+
426
+ def forward(
427
+ self,
428
+ hidden_states: torch.Tensor,
429
+ attention_mask: Optional[torch.Tensor] = None,
430
+ position_ids: Optional[torch.LongTensor] = None,
431
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
432
+ output_attentions: Optional[bool] = False,
433
+ use_cache: Optional[bool] = False,
434
+ **kwargs,
435
+ ) -> Tuple[
436
+ torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
437
+ ]:
438
+ """
439
+ Args:
440
+ hidden_states (`torch.FloatTensor`): 输入形状 `(batch, seq_len, embed_dim)`
441
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask 形状`(batch, sequence_length)`,
442
+ 填充使用0表示
443
+ output_attentions (`bool`, *optional*): 是否返回所有注意力层的注意力张量。
444
+ use_cache (`bool`, *optional*): 如果设置为 `True`,则返回 `past_key_values` 关键值状态,可用于加速解码
445
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): 缓存的之前kv状态
446
+ """
447
+
448
+ residual = hidden_states
449
+
450
+ hidden_states = self.input_layernorm(hidden_states)
451
+
452
+ # Self Attention
453
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
454
+ hidden_states=hidden_states,
455
+ attention_mask=attention_mask,
456
+ position_ids=position_ids,
457
+ past_key_value=past_key_value,
458
+ output_attentions=output_attentions,
459
+ use_cache=use_cache,
460
+ )
461
+ hidden_states = residual + hidden_states
462
+
463
+ # Fully Connected
464
+ residual = hidden_states
465
+ hidden_states = self.post_layernorm(hidden_states)
466
+ hidden_states = self.mlp(hidden_states)
467
+ hidden_states = residual + hidden_states
468
+
469
+ outputs = (hidden_states,)
470
+
471
+ if output_attentions:
472
+ outputs += (self_attn_weights,)
473
+
474
+ if use_cache:
475
+ outputs += (present_key_value,)
476
+
477
+ return outputs
478
+
479
+
480
+ class BuddyPreTrainedModel(PreTrainedModel):
481
+ config_class = BuddyGPTConfig
482
+ # 定义了模型内部子模块命名的基础前缀,当加载或保存模型时,这个前缀将用于识别模型主体部分。
483
+ base_model_prefix = "model"
484
+ # 表明该模型支持梯度检查点技术,这是一种内存优化策略,可减少模型训练时所需的显存
485
+ supports_gradient_checkpointing = True
486
+ # 指定了在序列化过程中不应被拆分的模块列表,即在模型保存与加载时保持这些模块作为一个整体。
487
+ _no_split_modules = ["DecoderLayer"]
488
+ # 在跨设备数据移动时,指示哪些关键字(key)对应的数据应该跳过设备放置步骤。
489
+ _skip_keys_device_placement = "past_key_values"
490
+ # Scaled Dot Product Attention (SDPA)
491
+ _supports_sdpa = True
492
+ # 表示模型支持缓存机制,这在自回归模型(如Transformer解码器)中很常见,
493
+ # 用于存储先前计算的结果以加快后续时间步长的计算速度。
494
+ _supports_cache_class = True
495
+
496
+ def _init_weights(self, module):
497
+ std = self.config.initializer_range
498
+ if isinstance(module, nn.Linear):
499
+ module.weight.data.normal_(mean=0.0, std=std)
500
+ if module.bias is not None:
501
+ module.bias.data.zero_()
502
+ elif isinstance(module, nn.Embedding):
503
+ module.weight.data.normal_(mean=0.0, std=std)
504
+ if module.padding_idx is not None:
505
+ module.weight.data[module.padding_idx].zero_()
506
+
507
+
508
+ class BuddyGPTModel(BuddyPreTrainedModel):
509
+ """根据配置文件堆叠 DecoderLayer
510
+ Args:
511
+ config: BuddyGPTConfig
512
+ """
513
+
514
+ def __init__(self, config: BuddyGPTConfig):
515
+ super().__init__(config)
516
+ self.padding_idx = config.pad_token_id
517
+ self.vocab_size = config.vocab_size
518
+
519
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
520
+ self.layers = nn.ModuleList([DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
521
+ self._attn_implementation = config._attn_implementation
522
+ self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
523
+
524
+ self.gradient_checkpointing = False
525
+ # Initialize weights and apply final processing
526
+ self.post_init()
527
+
528
+ def get_input_embeddings(self):
529
+ return self.embed_tokens
530
+
531
+ def set_input_embeddings(self, value):
532
+ self.embed_tokens = value
533
+
534
+ def forward(
535
+ self,
536
+ input_ids: torch.LongTensor = None,
537
+ attention_mask: Optional[torch.Tensor] = None,
538
+ position_ids: Optional[torch.LongTensor] = None, # 每个输入序列词元在位置嵌入中的位置索引
539
+ past_key_values: Optional[List[torch.FloatTensor]] = None, # 可用于加速序列解码预先计算的隐藏状态(自注意力块和交叉注意力块中的键和值)
540
+ inputs_embeds: Optional[torch.FloatTensor] = None,
541
+ use_cache: Optional[bool] = None,
542
+ output_attentions: Optional[bool] = None,
543
+ output_hidden_states: Optional[bool] = None,
544
+ return_dict: Optional[bool] = None,
545
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
546
+ output_attentions = (output_attentions if output_attentions is not None else self.config.output_attentions)
547
+ output_hidden_states = (output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states)
548
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
549
+
550
+ return_dict = (return_dict if return_dict is not None else self.config.use_return_dict)
551
+
552
+ # retrieve input_ids and inputs_embeds
553
+ if input_ids is not None and inputs_embeds is not None:
554
+ raise ValueError(
555
+ "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"
556
+ )
557
+ elif input_ids is not None:
558
+ batch_size, seq_length = input_ids.shape
559
+ elif inputs_embeds is not None:
560
+ batch_size, seq_length, _ = inputs_embeds.shape
561
+ else:
562
+ raise ValueError(
563
+ "You have to specify either decoder_input_ids or decoder_inputs_embeds"
564
+ )
565
+
566
+ if self.gradient_checkpointing and self.training:
567
+ if use_cache:
568
+ logger.warning_once(
569
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
570
+ )
571
+ use_cache = False
572
+
573
+ past_key_values_length = 0
574
+
575
+ if use_cache:
576
+ use_legacy_cache = not isinstance(past_key_values, Cache)
577
+ if use_legacy_cache:
578
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
579
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
580
+
581
+ if position_ids is None:
582
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
583
+ # 生成一个从past_key_values_length到seq_length + past_key_values_length的整数序列
584
+ position_ids = torch.arange(
585
+ past_key_values_length,
586
+ seq_length + past_key_values_length,
587
+ dtype=torch.long,
588
+ device=device,
589
+ )
590
+ # 将生成的序列重塑为形状为(1, seq_length)的张量,然后展平为形状为(-1, seq_length)的张量
591
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
592
+ else:
593
+ position_ids = position_ids.view(-1, seq_length).long()
594
+
595
+ if inputs_embeds is None:
596
+ inputs_embeds = self.embed_tokens(input_ids)
597
+
598
+ # 适应不同注意力机制对注意力掩码的不同要求而设计的
599
+ if self._attn_implementation == "sdpa" and not output_attentions:
600
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
601
+ # the manual implementation that requires a 4D causal mask in all cases.
602
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
603
+ attention_mask,
604
+ (batch_size, seq_length),
605
+ inputs_embeds,
606
+ past_key_values_length,
607
+ )
608
+ else:
609
+ # 4d mask is passed through the layers
610
+ attention_mask = _prepare_4d_causal_attention_mask(
611
+ attention_mask,
612
+ (batch_size, seq_length),
613
+ inputs_embeds,
614
+ past_key_values_length,
615
+ )
616
+
617
+ hidden_states = inputs_embeds
618
+
619
+ # decoder layers
620
+ all_hidden_states = () if output_hidden_states else None
621
+ all_self_attns = () if output_attentions else None
622
+ next_decoder_cache = None
623
+
624
+ for decoder_layer in self.layers:
625
+ # 1.隐藏状态保存
626
+ if output_hidden_states:
627
+ all_hidden_states += (hidden_states,)
628
+ # 2.梯度检查,方便在反向传播时只激活部分层,节省内存资源
629
+ # 3.解码层:
630
+ if self.gradient_checkpointing and self.training:
631
+ layer_outputs = self._gradient_checkpointing_func(
632
+ decoder_layer.__call__,
633
+ hidden_states,
634
+ attention_mask,
635
+ position_ids,
636
+ past_key_values,
637
+ output_attentions,
638
+ use_cache,
639
+ )
640
+ else:
641
+ layer_outputs = decoder_layer(
642
+ hidden_states,
643
+ attention_mask=attention_mask,
644
+ position_ids=position_ids,
645
+ past_key_value=past_key_values,
646
+ output_attentions=output_attentions,
647
+ use_cache=use_cache,
648
+ )
649
+ # 4.更新隐藏状态
650
+ hidden_states = layer_outputs[0]
651
+ # 5.更新缓存
652
+ if use_cache:
653
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
654
+ # 6.注意力输出保存
655
+ if output_attentions:
656
+ all_self_attns += (layer_outputs[1],)
657
+
658
+ hidden_states = self.norm(hidden_states)
659
+
660
+ # add hidden states from the last decoder layer
661
+ if output_hidden_states:
662
+ all_hidden_states += (hidden_states,)
663
+
664
+ next_cache = None
665
+ if use_cache:
666
+ next_cache = (
667
+ next_decoder_cache.to_legacy_cache()
668
+ if use_legacy_cache
669
+ else next_decoder_cache
670
+ )
671
+
672
+ if not return_dict:
673
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
674
+ return BaseModelOutputWithPast(
675
+ last_hidden_state=hidden_states,
676
+ past_key_values=next_cache,
677
+ hidden_states=all_hidden_states,
678
+ attentions=all_self_attns,
679
+ )
680
+
681
+
682
+ class BuddyGPTForCausalLM(BuddyPreTrainedModel, GenerationMixin):
683
+ _tied_weights_keys = ["lm_head.weight"]
684
+
685
+ def __init__(self, config):
686
+ super().__init__(config)
687
+ self.model = BuddyGPTModel(config)
688
+ self.vocab_size = config.vocab_size
689
+ self.tie_word_embeddings = config.tie_word_embeddings
690
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
691
+ if self.tie_word_embeddings:
692
+ self.lm_head.weight = self.model.embed_tokens.weight
693
+
694
+ # Initialize weights and apply final processing
695
+ self.post_init()
696
+
697
+ def get_input_embeddings(self):
698
+ return self.model.embed_tokens
699
+
700
+ def set_input_embeddings(self, value):
701
+ self.model.embed_tokens = value
702
+
703
+ def get_output_embeddings(self):
704
+ return self.lm_head
705
+
706
+ def set_output_embeddings(self, new_embeddings):
707
+ self.lm_head = new_embeddings
708
+
709
+ def set_decoder(self, decoder):
710
+ self.model = decoder
711
+
712
+ def get_decoder(self):
713
+ return self.model
714
+
715
+ def forward(
716
+ self,
717
+ input_ids: torch.LongTensor = None,
718
+ attention_mask: Optional[torch.Tensor] = None,
719
+ position_ids: Optional[torch.LongTensor] = None,
720
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
721
+ inputs_embeds: Optional[torch.FloatTensor] = None,
722
+ labels: Optional[torch.LongTensor] = None,
723
+ use_cache: Optional[bool] = None,
724
+ output_attentions: Optional[bool] = None,
725
+ output_hidden_states: Optional[bool] = None,
726
+ return_dict: Optional[bool] = None,
727
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
728
+
729
+ output_attentions = (output_attentions if output_attentions is not None else self.config.output_attentions)
730
+ output_hidden_states = (output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states)
731
+ return_dict = (return_dict if return_dict is not None else self.config.use_return_dict)
732
+
733
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
734
+ outputs = self.model(
735
+ input_ids=input_ids,
736
+ attention_mask=attention_mask,
737
+ position_ids=position_ids,
738
+ past_key_values=past_key_values,
739
+ inputs_embeds=inputs_embeds,
740
+ use_cache=use_cache,
741
+ output_attentions=output_attentions,
742
+ output_hidden_states=output_hidden_states,
743
+ return_dict=return_dict,
744
+ )
745
+
746
+ hidden_states = outputs[0]
747
+ logits = self.lm_head(hidden_states)
748
+ logits = logits.float()
749
+
750
+ loss = None
751
+ if labels is not None:
752
+ # Shift so that tokens < n predict n
753
+ # 对于自回归模型(如GPT系列),我们需要将模型输出的logits向前移动一位,
754
+ # 这样使得模型预测的是当前时刻 t 的下一个词,而非当前词本身
755
+ shift_logits = logits[..., :-1, :].contiguous()
756
+ # 同时,也需要将真实标签(labels)向前移动一位以与调整后的logits对齐
757
+ shift_labels = labels[..., 1:].contiguous()
758
+ # Flatten the tokens
759
+ loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
760
+
761
+ # 将移位后的 logits 和 labels 扁平化,即将它们展平为一维张量
762
+ # 其中shift_logits变成 (batch_size * sequence_length, vocab_size) 的形式
763
+ # shift_labels变为 (batch_size * sequence_length) 的形式
764
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
765
+ shift_labels = shift_labels.view(-1)
766
+
767
+ # Enable model parallelism
768
+ # 确保模型并行计算时,labels的数据存储位置与logits一致
769
+ shift_labels = shift_labels.to(shift_logits.device)
770
+ loss = loss_fct(shift_logits, shift_labels)
771
+
772
+ if not return_dict:
773
+ output = (logits,) + outputs[1:]
774
+ return (loss,) + output if loss is not None else output
775
+
776
+ return CausalLMOutputWithPast(
777
+ loss=loss,
778
+ logits=logits,
779
+ past_key_values=outputs.past_key_values,
780
+ hidden_states=outputs.hidden_states,
781
+ attentions=outputs.attentions,
782
+ )
783
+
784
+ def prepare_inputs_for_generation(
785
+ self,
786
+ input_ids,
787
+ past_key_values=None,
788
+ attention_mask=None,
789
+ inputs_embeds=None,
790
+ **kwargs,
791
+ ):
792
+ """准备模型的输入参数
793
+ 包括处理input_ids、past_key_values(历史隐藏状态缓存)、attention_mask以及可选的inputs_embeds。
794
+ """
795
+ # Omit tokens covered by past_key_values
796
+ if past_key_values is not None:
797
+ if isinstance(past_key_values, Cache):
798
+ cache_length = past_key_values.get_seq_length()
799
+ past_length = past_key_values.seen_tokens
800
+ max_cache_length = 2048
801
+ else:
802
+ cache_length = past_length = past_key_values[0][0].shape[2]
803
+ max_cache_length = None
804
+
805
+ # 根据缓存情况裁剪input_ids,只保留未处理的token:
806
+ # # 1. 如果 attention_mask 比 input_ids 更长,说明部分输入已通过缓存传递(如仅传入inputs_embeds)
807
+ if (
808
+ attention_mask is not None
809
+ and attention_mask.shape[1] > input_ids.shape[1]
810
+ ):
811
+ # 取最后未处理的部分
812
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
813
+ # 2. 若已处理的 token 数小于input_ids中的总数,表明input_ids包含全部输入,从中去掉已处理的部分
814
+ elif past_length < input_ids.shape[1]:
815
+ input_ids = input_ids[:, past_length:]
816
+ # 3. 否则,认为input_ids中只有待处理的新token
817
+
818
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
819
+ if (
820
+ max_cache_length is not None
821
+ and attention_mask is not None
822
+ and cache_length + input_ids.shape[1] > max_cache_length
823
+ ):
824
+ attention_mask = attention_mask[:, -max_cache_length:]
825
+
826
+ # 初始化或处理position_ids
827
+ position_ids = kwargs.get("position_ids", None)
828
+ # 如果attention_mask存在但position_ids不存在,则基于attention_mask动态创建position_ids
829
+ if attention_mask is not None and position_ids is None:
830
+ # create position_ids on the fly for batch generation
831
+ position_ids = attention_mask.long().cumsum(-1) - 1
832
+ position_ids.masked_fill_(attention_mask == 0, 1)
833
+ if past_key_values:
834
+ position_ids = position_ids[:, -input_ids.shape[1] :]
835
+
836
+ # 根据inputs_embeds和past_key_values的存在与否来决定模型输入
837
+ # 如果提供了inputs_embeds且没有past_key_values(首次生成步骤),则直接使用inputs_embeds作为模型输入
838
+ if inputs_embeds is not None and past_key_values is None:
839
+ model_inputs = {"inputs_embeds": inputs_embeds}
840
+ else:
841
+ model_inputs = {"input_ids": input_ids}
842
+
843
+ model_inputs.update(
844
+ {
845
+ "position_ids": position_ids,
846
+ "past_key_values": past_key_values,
847
+ "use_cache": kwargs.get("use_cache"),
848
+ "attention_mask": attention_mask,
849
+ }
850
+ )
851
+ return model_inputs
852
+
853
+ @staticmethod
854
+ def _reorder_cache(past_key_values, beam_idx):
855
+ """用于重新排序缓存中的历史隐藏状态,以适应束搜索(beam search)算法"""
856
+ reordered_past = ()
857
+ # 遍历每一层的隐藏状态
858
+ for layer_past in past_key_values:
859
+ # 对于每一层的每个隐藏状态向量,执行索引选择操作
860
+ reordered_past += (
861
+ tuple(
862
+ past_state.index_select(0, beam_idx.to(past_state.device))
863
+ for past_state in layer_past
864
+ ),
865
+ )
866
+ return reordered_past
867
+
868
+ def generate(
869
+ self,
870
+ inputs: Optional[torch.Tensor] = None,
871
+ generation_config: Optional[GenerationConfig] = None,
872
+ streamer = None,
873
+ **kwargs,
874
+ ):
875
+ if generation_config is None:
876
+ response = super().generate(
877
+ inputs,
878
+ generation_config=generation_config,
879
+ streamer=streamer,
880
+ **kwargs,
881
+ )
882
+
883
+ return response
884
+ repetition_penalty = kwargs.pop("repetition_penalty", generation_config.repetition_penalty)
885
+ generation_config.repetition_penalty = 1.0
886
+
887
+ logits_processor = None
888
+ if repetition_penalty > 1.0:
889
+ # warnings.warn("We highly recommend using OpenAI's frequency and presence penalty instead of the original repetition penalty. The original repetition penalty penalizes prompt tokens, which may lead to various potential issues. Therefore, your repetition penalty coefficient will be transformed into frequency penalty and presence penalty.", UserWarning)
890
+ presence_penalty = repetition_penalty - 1.0
891
+ frequency_penalty = repetition_penalty - 1.0
892
+ logits_processor = LogitsProcessorList(
893
+ [OutputRepetitionPenaltyLogitsProcessor(inputs.size(1), presence_penalty, frequency_penalty, 1.0)]
894
+ )
895
+
896
+ response = super().generate(
897
+ inputs,
898
+ generation_config=generation_config,
899
+ logits_processor=logits_processor,
900
+ streamer=streamer,
901
+ **kwargs,
902
+ )
903
+ generation_config.repetition_penalty = repetition_penalty
904
+ return response
905
+
906
+ def chat(
907
+ self,
908
+ tokenizer,
909
+ messages: List[dict],
910
+ system: str = "you are a helpful assistant!",
911
+ stream=False,
912
+ use_pot=False,
913
+ generation_config: Optional[GenerationConfig]=None
914
+ ):
915
+
916
+ generation_config = generation_config or self.generation_config
917
+ input_ids = make_context(
918
+ model=self, tokenizer=tokenizer, messages=messages,
919
+ system=system, max_new_tokens=generation_config.max_new_tokens
920
+ )
921
+
922
+ # for inputs in input_ids:
923
+ # print("decode: ", tokenizer.decode(inputs))
924
+
925
+ if stream:
926
+ streamer = TextIterStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, use_pot=use_pot)
927
+ Thread(target=self.generate, kwargs=dict(
928
+ inputs=input_ids, streamer=streamer,
929
+ generation_config=generation_config,
930
+ )).start()
931
+ return streamer
932
+ else:
933
+ generated_ids = self.generate(input_ids, generation_config=generation_config)
934
+ # response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)
935
+ generated_ids = [
936
+ output_ids[len(input_ids):] for input_ids, output_ids in zip(input_ids, generated_ids)
937
+ ]
938
+ response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
939
+ if use_pot:
940
+ response = parse_pot_no_stream(response)
941
+ return response
942
+
943
+ def print_model_parameters(model):
944
+ """打印模型各个层参数"""
945
+ param_sum = 0
946
+ for name, param in model.named_parameters():
947
+ if param.requires_grad:
948
+ param_sum += param.numel()
949
+ print(f"Layer: {name}, Parameters: {param.numel()}")
950
+ print(f"Total of parameters: {param_sum}")
951
+
952
+ # from transformers import AutoModel, AutoConfig, AutoModelForCausalLM
953
+
954
+ # AutoConfig.register("buddygpt", BuddyGPTConfig)
955
+ # AutoModel.register(BuddyGPTConfig, BuddyGPTModel)
956
+ # # AutoModelForCausalLM.register(BuddyGPTConfig, BuddyGPTModel)
957
+ from transformers.models.auto.configuration_auto import CONFIG_MAPPING
958
+ from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING
959
+
960
+ # 注册 config
961
+ CONFIG_MAPPING.register("buddygpt", BuddyGPTConfig)
962
+ # 注册模型
963
+ MODEL_FOR_CAUSAL_LM_MAPPING.register(BuddyGPTConfig, BuddyGPTForCausalLM)
964
+
965
+
966
+
special_tokens_map.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "additional_special_tokens": [
3
+ "<|im_start|>",
4
+ "<|im_end|>",
5
+ "<|object_ref_start|>",
6
+ "<|object_ref_end|>",
7
+ "<|box_start|>",
8
+ "<|box_end|>",
9
+ "<|quad_start|>",
10
+ "<|quad_end|>",
11
+ "<|vision_start|>",
12
+ "<|vision_end|>",
13
+ "<|vision_pad|>",
14
+ "<|image_pad|>",
15
+ "<|video_pad|>"
16
+ ],
17
+ "eos_token": {
18
+ "content": "<|im_end|>",
19
+ "lstrip": false,
20
+ "normalized": false,
21
+ "rstrip": false,
22
+ "single_word": false
23
+ },
24
+ "pad_token": {
25
+ "content": "<|endoftext|>",
26
+ "lstrip": false,
27
+ "normalized": false,
28
+ "rstrip": false,
29
+ "single_word": false
30
+ }
31
+ }
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4
3
+ size 11422654
tokenizer_config.json ADDED
@@ -0,0 +1,239 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_prefix_space": false,
4
+ "added_tokens_decoder": {
5
+ "151643": {
6
+ "content": "<|endoftext|>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ },
13
+ "151644": {
14
+ "content": "<|im_start|>",
15
+ "lstrip": false,
16
+ "normalized": false,
17
+ "rstrip": false,
18
+ "single_word": false,
19
+ "special": true
20
+ },
21
+ "151645": {
22
+ "content": "<|im_end|>",
23
+ "lstrip": false,
24
+ "normalized": false,
25
+ "rstrip": false,
26
+ "single_word": false,
27
+ "special": true
28
+ },
29
+ "151646": {
30
+ "content": "<|object_ref_start|>",
31
+ "lstrip": false,
32
+ "normalized": false,
33
+ "rstrip": false,
34
+ "single_word": false,
35
+ "special": true
36
+ },
37
+ "151647": {
38
+ "content": "<|object_ref_end|>",
39
+ "lstrip": false,
40
+ "normalized": false,
41
+ "rstrip": false,
42
+ "single_word": false,
43
+ "special": true
44
+ },
45
+ "151648": {
46
+ "content": "<|box_start|>",
47
+ "lstrip": false,
48
+ "normalized": false,
49
+ "rstrip": false,
50
+ "single_word": false,
51
+ "special": true
52
+ },
53
+ "151649": {
54
+ "content": "<|box_end|>",
55
+ "lstrip": false,
56
+ "normalized": false,
57
+ "rstrip": false,
58
+ "single_word": false,
59
+ "special": true
60
+ },
61
+ "151650": {
62
+ "content": "<|quad_start|>",
63
+ "lstrip": false,
64
+ "normalized": false,
65
+ "rstrip": false,
66
+ "single_word": false,
67
+ "special": true
68
+ },
69
+ "151651": {
70
+ "content": "<|quad_end|>",
71
+ "lstrip": false,
72
+ "normalized": false,
73
+ "rstrip": false,
74
+ "single_word": false,
75
+ "special": true
76
+ },
77
+ "151652": {
78
+ "content": "<|vision_start|>",
79
+ "lstrip": false,
80
+ "normalized": false,
81
+ "rstrip": false,
82
+ "single_word": false,
83
+ "special": true
84
+ },
85
+ "151653": {
86
+ "content": "<|vision_end|>",
87
+ "lstrip": false,
88
+ "normalized": false,
89
+ "rstrip": false,
90
+ "single_word": false,
91
+ "special": true
92
+ },
93
+ "151654": {
94
+ "content": "<|vision_pad|>",
95
+ "lstrip": false,
96
+ "normalized": false,
97
+ "rstrip": false,
98
+ "single_word": false,
99
+ "special": true
100
+ },
101
+ "151655": {
102
+ "content": "<|image_pad|>",
103
+ "lstrip": false,
104
+ "normalized": false,
105
+ "rstrip": false,
106
+ "single_word": false,
107
+ "special": true
108
+ },
109
+ "151656": {
110
+ "content": "<|video_pad|>",
111
+ "lstrip": false,
112
+ "normalized": false,
113
+ "rstrip": false,
114
+ "single_word": false,
115
+ "special": true
116
+ },
117
+ "151657": {
118
+ "content": "<tool_call>",
119
+ "lstrip": false,
120
+ "normalized": false,
121
+ "rstrip": false,
122
+ "single_word": false,
123
+ "special": false
124
+ },
125
+ "151658": {
126
+ "content": "</tool_call>",
127
+ "lstrip": false,
128
+ "normalized": false,
129
+ "rstrip": false,
130
+ "single_word": false,
131
+ "special": false
132
+ },
133
+ "151659": {
134
+ "content": "<|fim_prefix|>",
135
+ "lstrip": false,
136
+ "normalized": false,
137
+ "rstrip": false,
138
+ "single_word": false,
139
+ "special": false
140
+ },
141
+ "151660": {
142
+ "content": "<|fim_middle|>",
143
+ "lstrip": false,
144
+ "normalized": false,
145
+ "rstrip": false,
146
+ "single_word": false,
147
+ "special": false
148
+ },
149
+ "151661": {
150
+ "content": "<|fim_suffix|>",
151
+ "lstrip": false,
152
+ "normalized": false,
153
+ "rstrip": false,
154
+ "single_word": false,
155
+ "special": false
156
+ },
157
+ "151662": {
158
+ "content": "<|fim_pad|>",
159
+ "lstrip": false,
160
+ "normalized": false,
161
+ "rstrip": false,
162
+ "single_word": false,
163
+ "special": false
164
+ },
165
+ "151663": {
166
+ "content": "<|repo_name|>",
167
+ "lstrip": false,
168
+ "normalized": false,
169
+ "rstrip": false,
170
+ "single_word": false,
171
+ "special": false
172
+ },
173
+ "151664": {
174
+ "content": "<|file_sep|>",
175
+ "lstrip": false,
176
+ "normalized": false,
177
+ "rstrip": false,
178
+ "single_word": false,
179
+ "special": false
180
+ },
181
+ "151665": {
182
+ "content": "<tool_response>",
183
+ "lstrip": false,
184
+ "normalized": false,
185
+ "rstrip": false,
186
+ "single_word": false,
187
+ "special": false
188
+ },
189
+ "151666": {
190
+ "content": "</tool_response>",
191
+ "lstrip": false,
192
+ "normalized": false,
193
+ "rstrip": false,
194
+ "single_word": false,
195
+ "special": false
196
+ },
197
+ "151667": {
198
+ "content": "<think>",
199
+ "lstrip": false,
200
+ "normalized": false,
201
+ "rstrip": false,
202
+ "single_word": false,
203
+ "special": false
204
+ },
205
+ "151668": {
206
+ "content": "</think>",
207
+ "lstrip": false,
208
+ "normalized": false,
209
+ "rstrip": false,
210
+ "single_word": false,
211
+ "special": false
212
+ }
213
+ },
214
+ "additional_special_tokens": [
215
+ "<|im_start|>",
216
+ "<|im_end|>",
217
+ "<|object_ref_start|>",
218
+ "<|object_ref_end|>",
219
+ "<|box_start|>",
220
+ "<|box_end|>",
221
+ "<|quad_start|>",
222
+ "<|quad_end|>",
223
+ "<|vision_start|>",
224
+ "<|vision_end|>",
225
+ "<|vision_pad|>",
226
+ "<|image_pad|>",
227
+ "<|video_pad|>"
228
+ ],
229
+ "bos_token": null,
230
+ "clean_up_tokenization_spaces": false,
231
+ "eos_token": "<|im_end|>",
232
+ "errors": "replace",
233
+ "extra_special_tokens": {},
234
+ "model_max_length": 131072,
235
+ "pad_token": "<|endoftext|>",
236
+ "split_special_tokens": false,
237
+ "tokenizer_class": "Qwen2Tokenizer",
238
+ "unk_token": null
239
+ }
training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:29dac4e26f9eb136078d5b4f12f373b7ebe946996768daa65b738e30bbbe7720
3
+ size 5304
vocab.json ADDED
The diff for this file is too large to render. See raw diff