Upload 7 files
Browse files- config.json +35 -32
- generation_config.json +6 -4
- model_minimind.py +446 -0
- pytorch_model.bin +3 -0
- special_tokens_map.json +30 -0
- tokenizer.json +0 -0
- tokenizer_config.json +44 -43
config.json
CHANGED
|
@@ -1,32 +1,35 @@
|
|
| 1 |
-
{
|
| 2 |
-
"architectures": [
|
| 3 |
-
"
|
| 4 |
-
],
|
| 5 |
-
"auto_map": {
|
| 6 |
-
"AutoConfig": "
|
| 7 |
-
"AutoModelForCausalLM": "
|
| 8 |
-
},
|
| 9 |
-
"aux_loss_alpha": 0.1,
|
| 10 |
-
"
|
| 11 |
-
"dropout": 0.0,
|
| 12 |
-
"
|
| 13 |
-
"
|
| 14 |
-
"
|
| 15 |
-
"
|
| 16 |
-
"
|
| 17 |
-
"
|
| 18 |
-
"
|
| 19 |
-
"
|
| 20 |
-
"n_routed_experts": 4,
|
| 21 |
-
"n_shared_experts":
|
| 22 |
-
"
|
| 23 |
-
"
|
| 24 |
-
"num_experts_per_tok": 2,
|
| 25 |
-
"
|
| 26 |
-
"
|
| 27 |
-
"
|
| 28 |
-
"
|
| 29 |
-
"
|
| 30 |
-
"
|
| 31 |
-
"
|
| 32 |
-
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"architectures": [
|
| 3 |
+
"MiniMindForCausalLM"
|
| 4 |
+
],
|
| 5 |
+
"auto_map": {
|
| 6 |
+
"AutoConfig": "model_minimind.MiniMindConfig",
|
| 7 |
+
"AutoModelForCausalLM": "model_minimind.MiniMindForCausalLM"
|
| 8 |
+
},
|
| 9 |
+
"aux_loss_alpha": 0.1,
|
| 10 |
+
"bos_token_id": 1,
|
| 11 |
+
"dropout": 0.0,
|
| 12 |
+
"eos_token_id": 2,
|
| 13 |
+
"flash_attn": true,
|
| 14 |
+
"hidden_act": "silu",
|
| 15 |
+
"hidden_size": 640,
|
| 16 |
+
"intermediate_size": 1728,
|
| 17 |
+
"max_position_embeddings": 32768,
|
| 18 |
+
"max_seq_len": 8192,
|
| 19 |
+
"model_type": "minimind",
|
| 20 |
+
"n_routed_experts": 4,
|
| 21 |
+
"n_shared_experts": 1,
|
| 22 |
+
"norm_topk_prob": true,
|
| 23 |
+
"num_attention_heads": 8,
|
| 24 |
+
"num_experts_per_tok": 2,
|
| 25 |
+
"num_hidden_layers": 8,
|
| 26 |
+
"num_key_value_heads": 2,
|
| 27 |
+
"rms_norm_eps": 1e-05,
|
| 28 |
+
"rope_theta": 1000000.0,
|
| 29 |
+
"scoring_func": "softmax",
|
| 30 |
+
"seq_aux": true,
|
| 31 |
+
"torch_dtype": "bfloat16",
|
| 32 |
+
"transformers_version": "4.51.3",
|
| 33 |
+
"use_moe": true,
|
| 34 |
+
"vocab_size": 6400
|
| 35 |
+
}
|
generation_config.json
CHANGED
|
@@ -1,4 +1,6 @@
|
|
| 1 |
-
{
|
| 2 |
-
"_from_model_config": true,
|
| 3 |
-
"
|
| 4 |
-
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"_from_model_config": true,
|
| 3 |
+
"bos_token_id": 1,
|
| 4 |
+
"eos_token_id": 2,
|
| 5 |
+
"transformers_version": "4.51.3"
|
| 6 |
+
}
|
model_minimind.py
ADDED
|
@@ -0,0 +1,446 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
| 2 |
+
# MiniMind Config
|
| 3 |
+
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
| 4 |
+
|
| 5 |
+
from transformers import PretrainedConfig
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
class MiniMindConfig(PretrainedConfig):
|
| 9 |
+
model_type = "minimind"
|
| 10 |
+
|
| 11 |
+
def __init__(
|
| 12 |
+
self,
|
| 13 |
+
dropout: float = 0.0,
|
| 14 |
+
bos_token_id: int = 1,
|
| 15 |
+
eos_token_id: int = 2,
|
| 16 |
+
hidden_act: str = 'silu',
|
| 17 |
+
hidden_size: int = 512,
|
| 18 |
+
intermediate_size: int = None,
|
| 19 |
+
max_position_embeddings: int = 32768,
|
| 20 |
+
num_attention_heads: int = 8,
|
| 21 |
+
num_hidden_layers: int = 8,
|
| 22 |
+
num_key_value_heads: int = 2,
|
| 23 |
+
vocab_size: int = 6400,
|
| 24 |
+
rms_norm_eps: float = 1e-05,
|
| 25 |
+
rope_theta: int = 1000000.0,
|
| 26 |
+
flash_attn: bool = True,
|
| 27 |
+
####################################################
|
| 28 |
+
# Here are the specific configurations of MOE
|
| 29 |
+
# When use_moe is false, the following is invalid
|
| 30 |
+
####################################################
|
| 31 |
+
use_moe: bool = False,
|
| 32 |
+
num_experts_per_tok: int = 2,
|
| 33 |
+
n_routed_experts: int = 4,
|
| 34 |
+
n_shared_experts: int = 1,
|
| 35 |
+
scoring_func: str = 'softmax',
|
| 36 |
+
aux_loss_alpha: float = 0.1,
|
| 37 |
+
seq_aux: bool = True,
|
| 38 |
+
norm_topk_prob: bool = True,
|
| 39 |
+
**kwargs
|
| 40 |
+
):
|
| 41 |
+
super().__init__(**kwargs)
|
| 42 |
+
self.dropout = dropout
|
| 43 |
+
self.bos_token_id = bos_token_id
|
| 44 |
+
self.eos_token_id = eos_token_id
|
| 45 |
+
self.hidden_act = hidden_act
|
| 46 |
+
self.hidden_size = hidden_size
|
| 47 |
+
self.intermediate_size = intermediate_size
|
| 48 |
+
self.max_position_embeddings = max_position_embeddings
|
| 49 |
+
self.num_attention_heads = num_attention_heads
|
| 50 |
+
self.num_hidden_layers = num_hidden_layers
|
| 51 |
+
self.num_key_value_heads = num_key_value_heads
|
| 52 |
+
self.vocab_size = vocab_size
|
| 53 |
+
self.rms_norm_eps = rms_norm_eps
|
| 54 |
+
self.rope_theta = rope_theta
|
| 55 |
+
self.flash_attn = flash_attn
|
| 56 |
+
####################################################
|
| 57 |
+
# Here are the specific configurations of MOE
|
| 58 |
+
# When use_moe is false, the following is invalid
|
| 59 |
+
####################################################
|
| 60 |
+
self.use_moe = use_moe
|
| 61 |
+
self.num_experts_per_tok = num_experts_per_tok # 每个token选择的专家数量
|
| 62 |
+
self.n_routed_experts = n_routed_experts # 总的专家数量
|
| 63 |
+
self.n_shared_experts = n_shared_experts # 共享专家
|
| 64 |
+
self.scoring_func = scoring_func # 评分函数,默认为'softmax'
|
| 65 |
+
self.aux_loss_alpha = aux_loss_alpha # 辅助损失的alpha参数
|
| 66 |
+
self.seq_aux = seq_aux # 是否在序列级别上计算辅助损失
|
| 67 |
+
self.norm_topk_prob = norm_topk_prob # 是否标准化top-k概率
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
| 71 |
+
# MiniMind Model
|
| 72 |
+
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
| 73 |
+
|
| 74 |
+
import math
|
| 75 |
+
import torch
|
| 76 |
+
from torch import nn
|
| 77 |
+
from transformers.activations import ACT2FN
|
| 78 |
+
from typing import Optional, Tuple, List, Union
|
| 79 |
+
import torch.nn.functional as F
|
| 80 |
+
from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
|
| 81 |
+
from transformers.modeling_outputs import CausalLMOutputWithPast
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
class RMSNorm(torch.nn.Module):
|
| 85 |
+
def __init__(self, dim: int, eps: float = 1e-5):
|
| 86 |
+
super().__init__()
|
| 87 |
+
self.eps = eps
|
| 88 |
+
self.weight = nn.Parameter(torch.ones(dim))
|
| 89 |
+
|
| 90 |
+
def _norm(self, x):
|
| 91 |
+
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
| 92 |
+
|
| 93 |
+
def forward(self, x):
|
| 94 |
+
return self.weight * self._norm(x.float()).type_as(x)
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
|
| 98 |
+
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
|
| 99 |
+
t = torch.arange(end, device=freqs.device)
|
| 100 |
+
freqs = torch.outer(t, freqs).float()
|
| 101 |
+
freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
|
| 102 |
+
freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
|
| 103 |
+
return freqs_cos, freqs_sin
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
| 107 |
+
def rotate_half(x):
|
| 108 |
+
return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
|
| 109 |
+
|
| 110 |
+
q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
|
| 111 |
+
k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
|
| 112 |
+
return q_embed, k_embed
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
| 116 |
+
"""torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
|
| 117 |
+
bs, slen, num_key_value_heads, head_dim = x.shape
|
| 118 |
+
if n_rep == 1:
|
| 119 |
+
return x
|
| 120 |
+
return (
|
| 121 |
+
x[:, :, :, None, :]
|
| 122 |
+
.expand(bs, slen, num_key_value_heads, n_rep, head_dim)
|
| 123 |
+
.reshape(bs, slen, num_key_value_heads * n_rep, head_dim)
|
| 124 |
+
)
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
class Attention(nn.Module):
|
| 128 |
+
def __init__(self, args: MiniMindConfig):
|
| 129 |
+
super().__init__()
|
| 130 |
+
self.num_key_value_heads = args.num_attention_heads if args.num_key_value_heads is None else args.num_key_value_heads
|
| 131 |
+
assert args.num_attention_heads % self.num_key_value_heads == 0
|
| 132 |
+
self.n_local_heads = args.num_attention_heads
|
| 133 |
+
self.n_local_kv_heads = self.num_key_value_heads
|
| 134 |
+
self.n_rep = self.n_local_heads // self.n_local_kv_heads
|
| 135 |
+
self.head_dim = args.hidden_size // args.num_attention_heads
|
| 136 |
+
self.q_proj = nn.Linear(args.hidden_size, args.num_attention_heads * self.head_dim, bias=False)
|
| 137 |
+
self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
| 138 |
+
self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
| 139 |
+
self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False)
|
| 140 |
+
self.attn_dropout = nn.Dropout(args.dropout)
|
| 141 |
+
self.resid_dropout = nn.Dropout(args.dropout)
|
| 142 |
+
self.dropout = args.dropout
|
| 143 |
+
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
|
| 144 |
+
# print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
|
| 145 |
+
|
| 146 |
+
def forward(self,
|
| 147 |
+
x: torch.Tensor,
|
| 148 |
+
position_embeddings: Tuple[torch.Tensor, torch.Tensor], # 修改为接收cos和sin
|
| 149 |
+
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
| 150 |
+
use_cache=False,
|
| 151 |
+
attention_mask: Optional[torch.Tensor] = None):
|
| 152 |
+
bsz, seq_len, _ = x.shape
|
| 153 |
+
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
| 154 |
+
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
|
| 155 |
+
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
| 156 |
+
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
| 157 |
+
|
| 158 |
+
cos, sin = position_embeddings
|
| 159 |
+
xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len])
|
| 160 |
+
|
| 161 |
+
# kv_cache实现
|
| 162 |
+
if past_key_value is not None:
|
| 163 |
+
xk = torch.cat([past_key_value[0], xk], dim=1)
|
| 164 |
+
xv = torch.cat([past_key_value[1], xv], dim=1)
|
| 165 |
+
past_kv = (xk, xv) if use_cache else None
|
| 166 |
+
|
| 167 |
+
xq, xk, xv = (
|
| 168 |
+
xq.transpose(1, 2),
|
| 169 |
+
repeat_kv(xk, self.n_rep).transpose(1, 2),
|
| 170 |
+
repeat_kv(xv, self.n_rep).transpose(1, 2)
|
| 171 |
+
)
|
| 172 |
+
|
| 173 |
+
if False and self.flash and seq_len != 1:
|
| 174 |
+
dropout_p = self.dropout if self.training else 0.0
|
| 175 |
+
attn_mask = None
|
| 176 |
+
if attention_mask is not None:
|
| 177 |
+
attn_mask = attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, seq_len, -1)
|
| 178 |
+
attn_mask = attn_mask.bool() if attention_mask is not None else None
|
| 179 |
+
|
| 180 |
+
output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=True)
|
| 181 |
+
else:
|
| 182 |
+
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
|
| 183 |
+
scores = scores + torch.triu(
|
| 184 |
+
torch.full((seq_len, seq_len), float("-inf"), device=scores.device),
|
| 185 |
+
diagonal=1
|
| 186 |
+
).unsqueeze(0).unsqueeze(0) # scores+mask
|
| 187 |
+
|
| 188 |
+
if attention_mask is not None:
|
| 189 |
+
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
|
| 190 |
+
extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
|
| 191 |
+
scores = scores + extended_attention_mask
|
| 192 |
+
|
| 193 |
+
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
|
| 194 |
+
scores = self.attn_dropout(scores)
|
| 195 |
+
output = scores @ xv
|
| 196 |
+
|
| 197 |
+
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
|
| 198 |
+
output = self.resid_dropout(self.o_proj(output))
|
| 199 |
+
return output, past_kv
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
class FeedForward(nn.Module):
|
| 203 |
+
def __init__(self, config: MiniMindConfig):
|
| 204 |
+
super().__init__()
|
| 205 |
+
if config.intermediate_size is None:
|
| 206 |
+
intermediate_size = int(config.hidden_size * 8 / 3)
|
| 207 |
+
config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64)
|
| 208 |
+
self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 209 |
+
self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
| 210 |
+
self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
| 211 |
+
self.dropout = nn.Dropout(config.dropout)
|
| 212 |
+
self.act_fn = ACT2FN[config.hidden_act]
|
| 213 |
+
|
| 214 |
+
def forward(self, x):
|
| 215 |
+
return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
class MoEGate(nn.Module):
|
| 219 |
+
def __init__(self, config: MiniMindConfig):
|
| 220 |
+
super().__init__()
|
| 221 |
+
self.config = config
|
| 222 |
+
self.top_k = config.num_experts_per_tok
|
| 223 |
+
self.n_routed_experts = config.n_routed_experts
|
| 224 |
+
|
| 225 |
+
self.scoring_func = config.scoring_func
|
| 226 |
+
self.alpha = config.aux_loss_alpha
|
| 227 |
+
self.seq_aux = config.seq_aux
|
| 228 |
+
|
| 229 |
+
self.norm_topk_prob = config.norm_topk_prob
|
| 230 |
+
self.gating_dim = config.hidden_size
|
| 231 |
+
self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
|
| 232 |
+
self.reset_parameters()
|
| 233 |
+
|
| 234 |
+
def reset_parameters(self) -> None:
|
| 235 |
+
import torch.nn.init as init
|
| 236 |
+
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
| 237 |
+
|
| 238 |
+
def forward(self, hidden_states):
|
| 239 |
+
bsz, seq_len, h = hidden_states.shape
|
| 240 |
+
hidden_states = hidden_states.view(-1, h)
|
| 241 |
+
logits = F.linear(hidden_states, self.weight, None)
|
| 242 |
+
if self.scoring_func == 'softmax':
|
| 243 |
+
scores = logits.softmax(dim=-1)
|
| 244 |
+
else:
|
| 245 |
+
raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
|
| 246 |
+
|
| 247 |
+
topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
|
| 248 |
+
|
| 249 |
+
if self.top_k > 1 and self.norm_topk_prob:
|
| 250 |
+
denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
|
| 251 |
+
topk_weight = topk_weight / denominator
|
| 252 |
+
|
| 253 |
+
if self.training and self.alpha > 0.0:
|
| 254 |
+
scores_for_aux = scores
|
| 255 |
+
aux_topk = self.top_k
|
| 256 |
+
topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
|
| 257 |
+
if self.seq_aux:
|
| 258 |
+
scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
|
| 259 |
+
ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
|
| 260 |
+
ce.scatter_add_(1, topk_idx_for_aux_loss,
|
| 261 |
+
torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
|
| 262 |
+
seq_len * aux_topk / self.n_routed_experts)
|
| 263 |
+
aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
|
| 264 |
+
else:
|
| 265 |
+
mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
|
| 266 |
+
ce = mask_ce.float().mean(0)
|
| 267 |
+
Pi = scores_for_aux.mean(0)
|
| 268 |
+
fi = ce * self.n_routed_experts
|
| 269 |
+
aux_loss = (Pi * fi).sum() * self.alpha
|
| 270 |
+
else:
|
| 271 |
+
aux_loss = 0
|
| 272 |
+
return topk_idx, topk_weight, aux_loss
|
| 273 |
+
|
| 274 |
+
|
| 275 |
+
class MOEFeedForward(nn.Module):
|
| 276 |
+
def __init__(self, config: MiniMindConfig):
|
| 277 |
+
super().__init__()
|
| 278 |
+
self.config = config
|
| 279 |
+
self.experts = nn.ModuleList([
|
| 280 |
+
FeedForward(config)
|
| 281 |
+
for _ in range(config.n_routed_experts)
|
| 282 |
+
])
|
| 283 |
+
self.gate = MoEGate(config)
|
| 284 |
+
if config.n_shared_experts > 0:
|
| 285 |
+
self.shared_experts = nn.ModuleList([
|
| 286 |
+
FeedForward(config)
|
| 287 |
+
for _ in range(config.n_shared_experts)
|
| 288 |
+
])
|
| 289 |
+
|
| 290 |
+
def forward(self, x):
|
| 291 |
+
identity = x
|
| 292 |
+
orig_shape = x.shape
|
| 293 |
+
bsz, seq_len, _ = x.shape
|
| 294 |
+
# 使用门控机制选择专家
|
| 295 |
+
topk_idx, topk_weight, aux_loss = self.gate(x)
|
| 296 |
+
x = x.view(-1, x.shape[-1])
|
| 297 |
+
flat_topk_idx = topk_idx.view(-1)
|
| 298 |
+
if self.training:
|
| 299 |
+
x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
|
| 300 |
+
y = torch.empty_like(x, dtype=torch.float16)
|
| 301 |
+
for i, expert in enumerate(self.experts):
|
| 302 |
+
y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致
|
| 303 |
+
y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
|
| 304 |
+
y = y.view(*orig_shape)
|
| 305 |
+
else:
|
| 306 |
+
y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
|
| 307 |
+
if self.config.n_shared_experts > 0:
|
| 308 |
+
for expert in self.shared_experts:
|
| 309 |
+
y = y + expert(identity)
|
| 310 |
+
self.aux_loss = aux_loss
|
| 311 |
+
return y
|
| 312 |
+
|
| 313 |
+
@torch.no_grad()
|
| 314 |
+
def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
|
| 315 |
+
expert_cache = torch.zeros_like(x)
|
| 316 |
+
idxs = flat_expert_indices.argsort()
|
| 317 |
+
tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
|
| 318 |
+
token_idxs = idxs // self.config.num_experts_per_tok
|
| 319 |
+
# 当tokens_per_expert = [6, 15, 20, 26],tokens_per_expert.shape[0]即为专家数量(此时为4)
|
| 320 |
+
# 且token_idxs = [3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...] 时
|
| 321 |
+
# 意味token_idxs[:6] -> [3, 7, 19, 21, 24, 25]这6个位置属于专家0处理的token(每个token有可能被多个专家处理,这取决于num_experts_per_tok)
|
| 322 |
+
# 接下来9个位置token_idxs[6:15] -> [4, 5, 6, 10, 11, 12...]属于专家1处理的token...依此类推
|
| 323 |
+
for i, end_idx in enumerate(tokens_per_expert):
|
| 324 |
+
start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
|
| 325 |
+
if start_idx == end_idx:
|
| 326 |
+
continue
|
| 327 |
+
expert = self.experts[i]
|
| 328 |
+
exp_token_idx = token_idxs[start_idx:end_idx]
|
| 329 |
+
expert_tokens = x[exp_token_idx]
|
| 330 |
+
expert_out = expert(expert_tokens).to(expert_cache.dtype)
|
| 331 |
+
expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
|
| 332 |
+
expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)
|
| 333 |
+
|
| 334 |
+
return expert_cache
|
| 335 |
+
|
| 336 |
+
|
| 337 |
+
class MiniMindBlock(nn.Module):
|
| 338 |
+
def __init__(self, layer_id: int, config: MiniMindConfig):
|
| 339 |
+
super().__init__()
|
| 340 |
+
self.num_attention_heads = config.num_attention_heads
|
| 341 |
+
self.hidden_size = config.hidden_size
|
| 342 |
+
self.head_dim = config.hidden_size // config.num_attention_heads
|
| 343 |
+
self.self_attn = Attention(config)
|
| 344 |
+
|
| 345 |
+
self.layer_id = layer_id
|
| 346 |
+
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
| 347 |
+
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
| 348 |
+
self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
|
| 349 |
+
|
| 350 |
+
def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
|
| 351 |
+
residual = hidden_states
|
| 352 |
+
hidden_states, present_key_value = self.self_attn(
|
| 353 |
+
self.input_layernorm(hidden_states), position_embeddings,
|
| 354 |
+
past_key_value, use_cache, attention_mask
|
| 355 |
+
)
|
| 356 |
+
hidden_states += residual
|
| 357 |
+
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
|
| 358 |
+
return hidden_states, present_key_value
|
| 359 |
+
|
| 360 |
+
|
| 361 |
+
class MiniMindModel(nn.Module):
|
| 362 |
+
def __init__(self, config: MiniMindConfig):
|
| 363 |
+
super().__init__()
|
| 364 |
+
self.config = config
|
| 365 |
+
self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
|
| 366 |
+
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
| 367 |
+
self.dropout = nn.Dropout(config.dropout)
|
| 368 |
+
self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
|
| 369 |
+
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
| 370 |
+
|
| 371 |
+
freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.hidden_size // config.num_attention_heads,
|
| 372 |
+
end=config.max_position_embeddings, theta=config.rope_theta)
|
| 373 |
+
self.register_buffer("freqs_cos", freqs_cos, persistent=False)
|
| 374 |
+
self.register_buffer("freqs_sin", freqs_sin, persistent=False)
|
| 375 |
+
|
| 376 |
+
def forward(self,
|
| 377 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 378 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 379 |
+
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
| 380 |
+
use_cache: bool = False,
|
| 381 |
+
**kwargs):
|
| 382 |
+
batch_size, seq_length = input_ids.shape
|
| 383 |
+
past_key_values = past_key_values or [None] * len(self.layers)
|
| 384 |
+
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
|
| 385 |
+
|
| 386 |
+
hidden_states = self.dropout(self.embed_tokens(input_ids))
|
| 387 |
+
|
| 388 |
+
position_embeddings = (
|
| 389 |
+
self.freqs_cos[start_pos:start_pos + seq_length],
|
| 390 |
+
self.freqs_sin[start_pos:start_pos + seq_length]
|
| 391 |
+
)
|
| 392 |
+
|
| 393 |
+
presents = []
|
| 394 |
+
for layer_idx, (layer, past_key_value) in enumerate(zip(self.layers, past_key_values)):
|
| 395 |
+
hidden_states, present = layer(
|
| 396 |
+
hidden_states,
|
| 397 |
+
position_embeddings,
|
| 398 |
+
past_key_value=past_key_value,
|
| 399 |
+
use_cache=use_cache,
|
| 400 |
+
attention_mask=attention_mask
|
| 401 |
+
)
|
| 402 |
+
presents.append(present)
|
| 403 |
+
|
| 404 |
+
hidden_states = self.norm(hidden_states)
|
| 405 |
+
|
| 406 |
+
aux_loss = sum(
|
| 407 |
+
layer.mlp.aux_loss
|
| 408 |
+
for layer in self.layers
|
| 409 |
+
if isinstance(layer.mlp, MOEFeedForward)
|
| 410 |
+
)
|
| 411 |
+
|
| 412 |
+
return hidden_states, presents, aux_loss
|
| 413 |
+
|
| 414 |
+
|
| 415 |
+
class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
|
| 416 |
+
config_class = MiniMindConfig
|
| 417 |
+
|
| 418 |
+
def __init__(self, config: MiniMindConfig = None):
|
| 419 |
+
self.config = config or MiniMindConfig()
|
| 420 |
+
super().__init__(self.config)
|
| 421 |
+
self.model = MiniMindModel(self.config)
|
| 422 |
+
self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
|
| 423 |
+
self.model.embed_tokens.weight = self.lm_head.weight
|
| 424 |
+
self.OUT = CausalLMOutputWithPast()
|
| 425 |
+
|
| 426 |
+
def forward(self,
|
| 427 |
+
input_ids: Optional[torch.Tensor] = None,
|
| 428 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 429 |
+
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
| 430 |
+
use_cache: bool = False,
|
| 431 |
+
logits_to_keep: Union[int, torch.Tensor] = 0,
|
| 432 |
+
**args):
|
| 433 |
+
h, past_kvs, aux_loss = self.model(
|
| 434 |
+
input_ids=input_ids,
|
| 435 |
+
attention_mask=attention_mask,
|
| 436 |
+
past_key_values=past_key_values,
|
| 437 |
+
use_cache=use_cache,
|
| 438 |
+
**args
|
| 439 |
+
)
|
| 440 |
+
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
| 441 |
+
logits = self.lm_head(h[:, slice_indices, :])
|
| 442 |
+
self.OUT.__setitem__('last_hidden_state', h)
|
| 443 |
+
self.OUT.__setitem__('logits', logits)
|
| 444 |
+
self.OUT.__setitem__('aux_loss', aux_loss)
|
| 445 |
+
self.OUT.__setitem__('past_key_values', past_kvs)
|
| 446 |
+
return self.OUT
|
pytorch_model.bin
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:eda2077ba415f651ff07e3e1408d268e1dcdc15a86364454747753a5b12d9f0e
|
| 3 |
+
size 290118354
|
special_tokens_map.json
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token": {
|
| 3 |
+
"content": "<|im_start|>",
|
| 4 |
+
"lstrip": false,
|
| 5 |
+
"normalized": false,
|
| 6 |
+
"rstrip": false,
|
| 7 |
+
"single_word": false
|
| 8 |
+
},
|
| 9 |
+
"eos_token": {
|
| 10 |
+
"content": "<|im_end|>",
|
| 11 |
+
"lstrip": false,
|
| 12 |
+
"normalized": false,
|
| 13 |
+
"rstrip": false,
|
| 14 |
+
"single_word": false
|
| 15 |
+
},
|
| 16 |
+
"pad_token": {
|
| 17 |
+
"content": "<|endoftext|>",
|
| 18 |
+
"lstrip": false,
|
| 19 |
+
"normalized": false,
|
| 20 |
+
"rstrip": false,
|
| 21 |
+
"single_word": false
|
| 22 |
+
},
|
| 23 |
+
"unk_token": {
|
| 24 |
+
"content": "<|endoftext|>",
|
| 25 |
+
"lstrip": false,
|
| 26 |
+
"normalized": false,
|
| 27 |
+
"rstrip": false,
|
| 28 |
+
"single_word": false
|
| 29 |
+
}
|
| 30 |
+
}
|
tokenizer.json
CHANGED
|
The diff for this file is too large to render.
See raw diff
|
|
|
tokenizer_config.json
CHANGED
|
@@ -1,43 +1,44 @@
|
|
| 1 |
-
{
|
| 2 |
-
"add_bos_token": false,
|
| 3 |
-
"add_eos_token": false,
|
| 4 |
-
"add_prefix_space": false,
|
| 5 |
-
"added_tokens_decoder": {
|
| 6 |
-
"0": {
|
| 7 |
-
"content": "
|
| 8 |
-
"lstrip": false,
|
| 9 |
-
"normalized": false,
|
| 10 |
-
"rstrip": false,
|
| 11 |
-
"single_word": false,
|
| 12 |
-
"special": true
|
| 13 |
-
},
|
| 14 |
-
"1": {
|
| 15 |
-
"content": "
|
| 16 |
-
"lstrip": false,
|
| 17 |
-
"normalized": false,
|
| 18 |
-
"rstrip": false,
|
| 19 |
-
"single_word": false,
|
| 20 |
-
"special": true
|
| 21 |
-
},
|
| 22 |
-
"2": {
|
| 23 |
-
"content": "
|
| 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": "
|
| 33 |
-
"chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '
|
| 34 |
-
"clean_up_tokenization_spaces": false,
|
| 35 |
-
"eos_token": "
|
| 36 |
-
"
|
| 37 |
-
"
|
| 38 |
-
"
|
| 39 |
-
"
|
| 40 |
-
"
|
| 41 |
-
"
|
| 42 |
-
"
|
| 43 |
-
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"add_bos_token": false,
|
| 3 |
+
"add_eos_token": false,
|
| 4 |
+
"add_prefix_space": false,
|
| 5 |
+
"added_tokens_decoder": {
|
| 6 |
+
"0": {
|
| 7 |
+
"content": "<|endoftext|>",
|
| 8 |
+
"lstrip": false,
|
| 9 |
+
"normalized": false,
|
| 10 |
+
"rstrip": false,
|
| 11 |
+
"single_word": false,
|
| 12 |
+
"special": true
|
| 13 |
+
},
|
| 14 |
+
"1": {
|
| 15 |
+
"content": "<|im_start|>",
|
| 16 |
+
"lstrip": false,
|
| 17 |
+
"normalized": false,
|
| 18 |
+
"rstrip": false,
|
| 19 |
+
"single_word": false,
|
| 20 |
+
"special": true
|
| 21 |
+
},
|
| 22 |
+
"2": {
|
| 23 |
+
"content": "<|im_end|>",
|
| 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": "<|im_start|>",
|
| 33 |
+
"chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<|im_start|>system\\n' + system_message + '<|im_end|>\\n' }}{% else %}{{ '<|im_start|>system\\nYou are a helpful assistant<|im_end|>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|im_start|>user\\n' + content + '<|im_end|>\\n<|im_start|>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '<|im_end|>' + '\\n' }}{% endif %}{% endfor %}",
|
| 34 |
+
"clean_up_tokenization_spaces": false,
|
| 35 |
+
"eos_token": "<|im_end|>",
|
| 36 |
+
"extra_special_tokens": {},
|
| 37 |
+
"legacy": true,
|
| 38 |
+
"model_max_length": 32768,
|
| 39 |
+
"pad_token": "<|endoftext|>",
|
| 40 |
+
"sp_model_kwargs": {},
|
| 41 |
+
"spaces_between_special_tokens": false,
|
| 42 |
+
"tokenizer_class": "PreTrainedTokenizer",
|
| 43 |
+
"unk_token": "<|endoftext|>"
|
| 44 |
+
}
|