louishwh commited on
Commit
bf73816
·
verified ·
1 Parent(s): bf0056b

Upload model files

Browse files
config.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "Transformer"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "model.LMConfig",
7
+ "AutoModelForCausalLM": "model.Transformer"
8
+ },
9
+ "aux_loss_alpha": 0.01,
10
+ "dim": 512,
11
+ "dropout": 0.0,
12
+ "flash_attn": true,
13
+ "hidden_dim": null,
14
+ "max_seq_len": 512,
15
+ "model_type": "minimind",
16
+ "multiple_of": 64,
17
+ "n_heads": 16,
18
+ "n_kv_heads": 8,
19
+ "n_layers": 8,
20
+ "n_routed_experts": 4,
21
+ "n_shared_experts": true,
22
+ "norm_eps": 1e-05,
23
+ "norm_topk_prob": true,
24
+ "num_experts_per_tok": 2,
25
+ "scoring_func": "softmax",
26
+ "seq_aux": true,
27
+ "torch_dtype": "float32",
28
+ "transformers_version": "4.44.0",
29
+ "use_moe": false,
30
+ "vocab_size": 6400
31
+ }
generation_config.json ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "transformers_version": "4.44.0"
4
+ }
model.py ADDED
@@ -0,0 +1,527 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # !/usr/bin/env python
2
+ # -*-coding:utf-8 -*-
3
+
4
+ """
5
+ # Time :2024/10/16 19:28
6
+ # Author :Maxwell
7
+ # Description:
8
+ """
9
+ import math
10
+ from typing import Any, Optional, Tuple
11
+ import torch
12
+ import torch.nn.functional as F
13
+ from torch import nn
14
+ from transformers import PreTrainedModel
15
+ from transformers.modeling_outputs import CausalLMOutputWithPast
16
+ from transformers import PretrainedConfig
17
+
18
+
19
+ class LMConfig(PretrainedConfig):
20
+ model_type = "minimind"
21
+
22
+ def __init__(
23
+ self,
24
+ dim: int = 512, # 模型的隐藏维度,即每个层的输出向量的大小
25
+ n_layers: int = 8, # Transformer 模型中的层数
26
+ n_heads: int = 16, # 多头注意力机制中的注意力头数
27
+ n_kv_heads: int = 8, # 键和值所使用的注意力头数
28
+ vocab_size: int = 6400, # 词汇表的大小
29
+ hidden_dim: int = None, # 前馈神经网络层的隐藏层维度
30
+ multiple_of: int = 64, # 用于计算隐藏层维度的倍数
31
+ norm_eps: float = 1e-5, # 归一化层中的 epsilon 值
32
+ max_seq_len: int = 512, # 模型支持的最大序列长度
33
+ dropout: float = 0.0, # dropout 概率,用于防止过拟合
34
+ flash_attn: bool = True, # 是否使用闪存注意力
35
+
36
+ # 下面是门控专家(MoE)特定的参数
37
+ use_moe: bool = False, # 是否启用门控专家机制(MoE)
38
+ num_experts_per_tok=2, # 每个标记选择的专家数量
39
+ n_routed_experts=4, # 总的专家数量
40
+ n_shared_experts: bool = True, # 是否使用共享专家
41
+ scoring_func='softmax', # 用于选择专家的评分函数
42
+ aux_loss_alpha=0.01, # 辅助损失的 alpha 参数
43
+ seq_aux=True, # 是否在序列级别上计算辅助损失
44
+ norm_topk_prob=True, # 是否标准化 top-k 概率
45
+ **kwargs,
46
+ ):
47
+ self.dim = dim
48
+ self.n_layers = n_layers
49
+ self.n_heads = n_heads
50
+ self.n_kv_heads = n_kv_heads
51
+ self.vocab_size = vocab_size
52
+ self.hidden_dim = hidden_dim
53
+ self.multiple_of = multiple_of
54
+ self.norm_eps = norm_eps
55
+ self.max_seq_len = max_seq_len
56
+ self.dropout = dropout
57
+ self.flash_attn = flash_attn
58
+
59
+ # 下面是门控专家(MoE)特定的配置
60
+ self.use_moe = use_moe
61
+ self.num_experts_per_tok = num_experts_per_tok
62
+ self.n_routed_experts = n_routed_experts
63
+ self.n_shared_experts = n_shared_experts
64
+ self.scoring_func = scoring_func
65
+ self.aux_loss_alpha = aux_loss_alpha
66
+ self.seq_aux = seq_aux
67
+ self.norm_topk_prob = norm_topk_prob
68
+ super().__init__(**kwargs)
69
+
70
+
71
+ class RMSNorm(torch.nn.Module):
72
+
73
+ """
74
+ 实现了 RMSNorm(Root Mean Square Normalization),用于对输入进行归一化。
75
+ 计算输入的均方根,并进行标准化处理。
76
+ """
77
+
78
+ def __init__(self, dim: int, eps: float):
79
+ super().__init__()
80
+ self.eps = eps
81
+ self.weight = nn.Parameter(torch.ones(dim))
82
+
83
+ def _norm(self, x):
84
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
85
+
86
+ def forward(self, x):
87
+ output = self._norm(x.float()).type_as(x)
88
+ return output * self.weight
89
+
90
+
91
+ def precompute_pos_cis(dim: int, end: int, theta: float = 10000.0, train_len: int = 512):
92
+ """
93
+ 预计算位置编码,使用复数形式表示位置编码以支持旋转嵌入(rotary embeddings)。
94
+ """
95
+ freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
96
+ t = torch.arange(end, device=freqs.device) # type: ignore
97
+ freqs = torch.outer(t, freqs).float() # type: ignore
98
+ pos_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
99
+ # 计算缩放因子
100
+ # scale = train_len / end
101
+ # 缩放旋转嵌入,实现线性的长度外推(注释掉不用是因为小模型依赖pos_cis拟合严重,直接做线性外推效果并不好)
102
+ # pos_cis = pos_cis * scale
103
+ return pos_cis
104
+
105
+
106
+ def apply_rotary_emb(xq, xk, pos_cis):
107
+ """
108
+ 应用旋转嵌入,通过复数运算实现位置编码的旋转。
109
+ :param xq:
110
+ :param xk:
111
+ :param pos_cis:
112
+ :return:
113
+ """
114
+ def unite_shape(pos_cis, x):
115
+ ndim = x.ndim
116
+ assert 0 <= 1 < ndim
117
+ assert pos_cis.shape == (x.shape[1], x.shape[-1])
118
+ shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
119
+ return pos_cis.view(*shape)
120
+
121
+ xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
122
+ xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
123
+ pos_cis = unite_shape(pos_cis, xq_)
124
+ xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
125
+ xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
126
+ return xq_out.type_as(xq), xk_out.type_as(xk)
127
+
128
+
129
+ def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
130
+ """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
131
+ bs, slen, n_kv_heads, head_dim = x.shape
132
+ if n_rep == 1:
133
+ return x
134
+ return (
135
+ x[:, :, :, None, :]
136
+ .expand(bs, slen, n_kv_heads, n_rep, head_dim)
137
+ .reshape(bs, slen, n_kv_heads * n_rep, head_dim)
138
+ )
139
+
140
+
141
+ class Attention(nn.Module):
142
+
143
+ """
144
+ 实现多头注意力机制。
145
+ 包括查询(query)、键(key)、值(value)线性变换,掩码机制,注意力得分计算等。
146
+ """
147
+
148
+ def __init__(self, args: LMConfig):
149
+ super().__init__()
150
+ self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
151
+ assert args.n_heads % self.n_kv_heads == 0
152
+ self.n_local_heads = args.n_heads
153
+ self.n_local_kv_heads = self.n_kv_heads
154
+ self.n_rep = self.n_local_heads // self.n_local_kv_heads
155
+ self.head_dim = args.dim // args.n_heads
156
+ self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
157
+ self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
158
+ self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
159
+ self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
160
+ self.k_cache, self.v_cache = None, None
161
+ self.attn_dropout = nn.Dropout(args.dropout)
162
+ self.resid_dropout = nn.Dropout(args.dropout)
163
+ self.dropout = args.dropout
164
+ self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
165
+
166
+ # print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
167
+ mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
168
+ mask = torch.triu(mask, diagonal=1)
169
+ self.register_buffer("mask", mask, persistent=False)
170
+
171
+ def forward(self, x: torch.Tensor, pos_cis: torch.Tensor, kv_cache=False):
172
+ bsz, seqlen, _ = x.shape
173
+
174
+ xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
175
+
176
+ xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
177
+ xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
178
+ xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
179
+
180
+ xq, xk = apply_rotary_emb(xq, xk, pos_cis)
181
+
182
+ # 更高效的kv_cache实现
183
+ if kv_cache and self.eval():
184
+ if seqlen == 1 and all(cache is not None for cache in (self.k_cache, self.v_cache)):
185
+ xk = torch.cat((self.k_cache, xk), dim=1)
186
+ xv = torch.cat((self.v_cache, xv), dim=1)
187
+ self.k_cache, self.v_cache = xk, xv
188
+
189
+ xk = repeat_kv(xk, self.n_rep) # (bs, seqlen, n_local_heads, head_dim)
190
+ xv = repeat_kv(xv, self.n_rep) # (bs, seqlen, n_local_heads, head_dim)
191
+
192
+ xq = xq.transpose(1, 2)
193
+ xk = xk.transpose(1, 2)
194
+ xv = xv.transpose(1, 2)
195
+
196
+ if self.flash and seqlen != 1:
197
+ output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None,
198
+ dropout_p=self.dropout if self.training else 0.0,
199
+ is_causal=True)
200
+ else:
201
+ scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
202
+ scores = scores + self.mask[:, :, :seqlen, :seqlen] # (bs, n_local_heads, seqlen, cache_len + seqlen)
203
+ scores = F.softmax(scores.float(), dim=-1).type_as(xq)
204
+ scores = self.attn_dropout(scores)
205
+ output = torch.matmul(scores, xv) # (bs, n_local_heads, seqlen, head_dim)
206
+
207
+ output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
208
+
209
+ output = self.wo(output)
210
+ output = self.resid_dropout(output)
211
+ return output
212
+
213
+
214
+ class FeedForward(nn.Module):
215
+ """
216
+ 实现前馈神经网络层,包含两个线性变换和一个激活函数(SiLU)。
217
+ """
218
+ def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):
219
+ super().__init__()
220
+ if hidden_dim is None:
221
+ hidden_dim = 4 * dim
222
+ hidden_dim = int(2 * hidden_dim / 3)
223
+ hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
224
+ self.w1 = nn.Linear(dim, hidden_dim, bias=False)
225
+ self.w2 = nn.Linear(hidden_dim, dim, bias=False)
226
+ self.w3 = nn.Linear(dim, hidden_dim, bias=False)
227
+ self.dropout = nn.Dropout(dropout)
228
+
229
+ def forward(self, x):
230
+ return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
231
+
232
+
233
+ class MoEGate(nn.Module):
234
+
235
+ """
236
+ 实现门控专家机制(Mixture of Experts),根据输入选择不同的专家网络进行计算。
237
+ 包含门控机制用于选择专家,和计算辅助损失(auxiliary loss)。
238
+ """
239
+ def __init__(self, config: LMConfig):
240
+ super().__init__()
241
+ self.config = config
242
+ self.top_k = config.num_experts_per_tok
243
+ self.n_routed_experts = config.n_routed_experts
244
+
245
+ self.scoring_func = config.scoring_func
246
+ self.alpha = config.aux_loss_alpha
247
+ self.seq_aux = config.seq_aux
248
+
249
+ self.norm_topk_prob = config.norm_topk_prob
250
+ self.gating_dim = config.dim
251
+ self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
252
+ self.reset_parameters()
253
+
254
+ def reset_parameters(self) -> None:
255
+ import torch.nn.init as init
256
+ init.kaiming_uniform_(self.weight, a=math.sqrt(5))
257
+
258
+ def forward(self, hidden_states):
259
+ bsz, seq_len, h = hidden_states.shape
260
+
261
+ hidden_states = hidden_states.view(-1, h)
262
+ logits = F.linear(hidden_states, self.weight, None)
263
+ if self.scoring_func == 'softmax':
264
+ scores = logits.softmax(dim=-1)
265
+ else:
266
+ raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
267
+
268
+ topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
269
+
270
+ if self.top_k > 1 and self.norm_topk_prob:
271
+ denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
272
+ topk_weight = topk_weight / denominator
273
+
274
+ if self.training and self.alpha > 0.0:
275
+ scores_for_aux = scores
276
+ aux_topk = self.top_k
277
+ topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
278
+ if self.seq_aux:
279
+ scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
280
+ ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
281
+ ce.scatter_add_(1, topk_idx_for_aux_loss,
282
+ torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
283
+ seq_len * aux_topk / self.n_routed_experts)
284
+ aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
285
+ else:
286
+ mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
287
+ ce = mask_ce.float().mean(0)
288
+ Pi = scores_for_aux.mean(0)
289
+ fi = ce * self.n_routed_experts
290
+ aux_loss = (Pi * fi).sum() * self.alpha
291
+ else:
292
+ aux_loss = None
293
+ return topk_idx, topk_weight, aux_loss
294
+
295
+
296
+ class MOEFeedForward(nn.Module):
297
+
298
+ def __init__(self, config: LMConfig):
299
+ super().__init__()
300
+ self.config = config
301
+ self.experts = nn.ModuleList([
302
+ FeedForward(
303
+ dim=config.dim,
304
+ hidden_dim=config.hidden_dim,
305
+ multiple_of=config.multiple_of,
306
+ dropout=config.dropout,
307
+ )
308
+ for _ in range(config.n_routed_experts)
309
+ ])
310
+
311
+ self.gate = MoEGate(config)
312
+ if config.n_shared_experts is not None:
313
+ self.shared_experts = FeedForward(
314
+ dim=config.dim,
315
+ hidden_dim=config.hidden_dim,
316
+ multiple_of=config.multiple_of,
317
+ dropout=config.dropout,
318
+ )
319
+
320
+ def forward(self, x):
321
+ identity = x
322
+ orig_shape = x.shape
323
+ bsz, seq_len, _ = x.shape
324
+
325
+ # 使用门控机制选择专家
326
+ topk_idx, topk_weight, aux_loss = self.gate(x)
327
+
328
+ x = x.view(-1, x.shape[-1])
329
+ flat_topk_idx = topk_idx.view(-1)
330
+
331
+ if self.training:
332
+ # 训练模式下,重复输入数据
333
+ x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
334
+ y = torch.empty_like(x, dtype=torch.float16)
335
+ for i, expert in enumerate(self.experts):
336
+ y[flat_topk_idx == i] = expert(x[flat_topk_idx == i])
337
+ y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
338
+ y = y.view(*orig_shape)
339
+ else:
340
+ # 推理模式下,只选择最优专家
341
+ y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
342
+
343
+ if self.config.n_shared_experts is not None:
344
+ y = y + self.shared_experts(identity)
345
+
346
+ return y
347
+
348
+ @torch.no_grad()
349
+ def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
350
+ expert_cache = torch.zeros_like(x)
351
+ idxs = flat_expert_indices.argsort()
352
+ tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
353
+ token_idxs = idxs // self.config.num_experts_per_tok
354
+ # 例如当tokens_per_expert=[6, 15, 20, 26, 33, 38, 46, 52]
355
+ # 当token_idxs=[3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...]
356
+ # 意味着当token_idxs[:6] -> [3, 7, 19, 21, 24, 25, 4]位置的token都由专家0处理,token_idxs[6:15]位置的token都由专家1处理......
357
+ for i, end_idx in enumerate(tokens_per_expert):
358
+ start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
359
+ if start_idx == end_idx:
360
+ continue
361
+ expert = self.experts[i]
362
+ exp_token_idx = token_idxs[start_idx:end_idx]
363
+ expert_tokens = x[exp_token_idx]
364
+ expert_out = expert(expert_tokens)
365
+ expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
366
+ # 使用 scatter_add_ 进行 sum 操作
367
+ expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)
368
+
369
+ return expert_cache
370
+
371
+
372
+ class TransformerBlock(nn.Module):
373
+
374
+ """
375
+ 组合注意力层和前馈层,构成 Transformer 的基本构建块。
376
+ 使用 RMSNorm 进行层归一化。
377
+ """
378
+ def __init__(self, layer_id: int, args: LMConfig):
379
+ super().__init__()
380
+ self.n_heads = args.n_heads
381
+ self.dim = args.dim
382
+ self.head_dim = args.dim // args.n_heads
383
+ self.attention = Attention(args)
384
+
385
+ self.layer_id = layer_id
386
+ self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
387
+ self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
388
+
389
+ if args.use_moe:
390
+ self.feed_forward = MOEFeedForward(args)
391
+ else:
392
+ self.feed_forward = FeedForward(
393
+ dim=args.dim,
394
+ hidden_dim=args.hidden_dim,
395
+ multiple_of=args.multiple_of,
396
+ dropout=args.dropout,
397
+ )
398
+
399
+ def forward(self, x, pos_cis, kv_cache=False):
400
+ h = x + self.attention(self.attention_norm(x), pos_cis, kv_cache)
401
+ out = h + self.feed_forward(self.ffn_norm(h))
402
+ return out
403
+
404
+
405
+ class Transformer(PreTrainedModel):
406
+ """
407
+ 主模型类,继承自 PreTrainedModel。
408
+ 包含嵌入层、多个 Transformer 块、输出层。
409
+ 支持生成文本、评估答案、处理输入输出。
410
+ """
411
+
412
+ config_class = LMConfig
413
+ last_loss: Optional[torch.Tensor]
414
+
415
+ def __init__(self, params: LMConfig = None):
416
+ super().__init__(params)
417
+ if not params:
418
+ params = LMConfig()
419
+ self.params = params
420
+ self.vocab_size = params.vocab_size
421
+ self.n_layers = params.n_layers
422
+
423
+ self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
424
+ self.dropout = nn.Dropout(params.dropout)
425
+ self.layers = torch.nn.ModuleList()
426
+ for layer_id in range(self.n_layers):
427
+ self.layers.append(TransformerBlock(layer_id, params))
428
+ self.norm = RMSNorm(params.dim, eps=params.norm_eps)
429
+ self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
430
+ self.tok_embeddings.weight = self.output.weight
431
+ pos_cis = precompute_pos_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len)
432
+ self.register_buffer("pos_cis", pos_cis, persistent=False)
433
+
434
+ self.apply(self._init_weights)
435
+
436
+ for pn, p in self.named_parameters():
437
+ if pn.endswith('w3.weight') or pn.endswith('wo.weight'):
438
+ torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * params.n_layers))
439
+
440
+ self.last_loss = None
441
+ self.OUT = CausalLMOutputWithPast()
442
+ self._no_split_modules = [name for name, _ in self.named_modules()]
443
+
444
+ def _init_weights(self, module):
445
+ if isinstance(module, nn.Linear):
446
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
447
+ if module.bias is not None:
448
+ torch.nn.init.zeros_(module.bias)
449
+ elif isinstance(module, nn.Embedding):
450
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
451
+
452
+ def forward(self, tokens: Optional[torch.Tensor] = None, targets: Optional[torch.Tensor] = None,
453
+ kv_cache=False, **keyargs):
454
+ current_idx = 0
455
+ if 'input_ids' in keyargs:
456
+ tokens = keyargs['input_ids']
457
+ if 'attention_mask' in keyargs:
458
+ targets = keyargs['attention_mask']
459
+ if 'current_idx' in keyargs:
460
+ current_idx = int(keyargs['current_idx'])
461
+
462
+ _bsz, seqlen = tokens.shape
463
+ h = self.tok_embeddings(tokens)
464
+ h = self.dropout(h)
465
+ pos_cis = self.pos_cis[current_idx:current_idx + seqlen]
466
+ for idx, layer in enumerate(self.layers):
467
+ h = layer(h, pos_cis, kv_cache)
468
+
469
+ h = self.norm(h)
470
+
471
+ if targets is not None:
472
+ logits = self.output(h)
473
+ self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1),
474
+ ignore_index=0, reduction='none')
475
+ else:
476
+ logits = self.output(h[:, [-1], :])
477
+ self.last_loss = None
478
+
479
+ self.OUT.__setitem__('logits', logits)
480
+ self.OUT.__setitem__('last_loss', self.last_loss)
481
+ return self.OUT
482
+
483
+ @torch.inference_mode()
484
+ def generate(self, idx, eos, max_new_tokens, temperature=0.7, top_k=8, stream=True, rp=1., kv_cache=True):
485
+ # rp: repetition_penalty
486
+ index = idx.shape[1]
487
+ init_inference = True
488
+ while idx.shape[1] < max_new_tokens - 1:
489
+ if init_inference or not kv_cache:
490
+ inference_res, init_inference = self(idx, kv_cache=kv_cache), False
491
+ else:
492
+ inference_res = self(idx[:, -1:], kv_cache=kv_cache, current_idx=idx.shape[1] - 1)
493
+
494
+ logits = inference_res.logits
495
+ logits = logits[:, -1, :]
496
+
497
+ for token in set(idx.tolist()[0]):
498
+ logits[:, token] /= rp
499
+
500
+ if temperature == 0.0:
501
+ _, idx_next = torch.topk(logits, k=1, dim=-1)
502
+ else:
503
+ logits = logits / temperature
504
+ if top_k is not None:
505
+ v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
506
+ logits[logits < v[:, [-1]]] = -float('Inf')
507
+
508
+ probs = F.softmax(logits, dim=-1)
509
+ idx_next = torch.multinomial(probs, num_samples=1, generator=None)
510
+
511
+ if idx_next == eos:
512
+ break
513
+
514
+ idx = torch.cat((idx, idx_next), dim=1)
515
+ if stream:
516
+ yield idx[:, index:]
517
+
518
+ if not stream:
519
+ yield idx[:, index:]
520
+
521
+ @torch.inference_mode()
522
+ def eval_answer(self, idx):
523
+ idx_cond = idx if idx.size(1) <= self.params.max_seq_len else idx[:, -self.params.max_seq_len:]
524
+ inference_res = self(idx_cond)
525
+ logits = inference_res.logits
526
+ logits = logits[:, -1, :]
527
+ return logits
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:676fb91edbba6e3e4b021e07a79a322f24f940352d300382ab4854465678ea64
3
+ size 107537292
special_tokens_map.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": {
3
+ "content": "<s>",
4
+ "lstrip": false,
5
+ "normalized": false,
6
+ "rstrip": false,
7
+ "single_word": false
8
+ },
9
+ "eos_token": {
10
+ "content": "</s>",
11
+ "lstrip": false,
12
+ "normalized": false,
13
+ "rstrip": false,
14
+ "single_word": false
15
+ },
16
+ "unk_token": {
17
+ "content": "<unk>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ }
23
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_eos_token": false,
4
+ "add_prefix_space": true,
5
+ "added_tokens_decoder": {
6
+ "0": {
7
+ "content": "<unk>",
8
+ "lstrip": false,
9
+ "normalized": false,
10
+ "rstrip": false,
11
+ "single_word": false,
12
+ "special": true
13
+ },
14
+ "1": {
15
+ "content": "<s>",
16
+ "lstrip": false,
17
+ "normalized": false,
18
+ "rstrip": false,
19
+ "single_word": false,
20
+ "special": true
21
+ },
22
+ "2": {
23
+ "content": "</s>",
24
+ "lstrip": false,
25
+ "normalized": false,
26
+ "rstrip": false,
27
+ "single_word": false,
28
+ "special": true
29
+ }
30
+ },
31
+ "additional_special_tokens": [],
32
+ "bos_token": "<s>",
33
+ "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{% endif %}{% if system_message is defined %}{{ system_message }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<s>user\\n' + content + '</s>\\n<s>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '</s>' + '\\n' }}{% endif %}{% endfor %}",
34
+ "clean_up_tokenization_spaces": false,
35
+ "eos_token": "</s>",
36
+ "legacy": true,
37
+ "model_max_length": 1000000000000000019884624838656,
38
+ "pad_token": null,
39
+ "sp_model_kwargs": {},
40
+ "spaces_between_special_tokens": false,
41
+ "tokenizer_class": "PreTrainedTokenizerFast",
42
+ "unk_token": "<unk>",
43
+ "use_default_system_prompt": false
44
+ }