Text Generation
Transformers
PyTorch
tinymixtral
conversational
custom_code
mikecovlee commited on
Commit
afb3d19
·
verified ·
1 Parent(s): b671dce

Upload 12 files

Browse files
LICENSE ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Michael Lee (李登淳)
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
README.md CHANGED
@@ -1,3 +1,105 @@
1
- ---
2
- license: mit
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # TinyMixtral 1B — Post-Trained
2
+
3
+ A 1.18B-parameter Mixture-of-Experts language model (351M active), post-trained on 1B tokens of educational and web text.
4
+
5
+ ## Model Details
6
+
7
+ | Property | Value |
8
+ |---|---|
9
+ | Architecture | Decoder-only Transformer with Sparse MoE |
10
+ | Total Parameters | 1,182,172,160 |
11
+ | Active Parameters | ~351M |
12
+ | Hidden Size | 1024 |
13
+ | Layers | 16 |
14
+ | Experts | 8 (top-2 routing) |
15
+ | Attention Heads | 16 query / 4 key-value (GQA) |
16
+ | Head Dimension | 64 |
17
+ | Intermediate Size | 2,816 (per expert) |
18
+ | Vocabulary | 32,000 |
19
+ | Context Length | 2,048 |
20
+ | Position Encoding | RoPE (theta=1e6) |
21
+ | Activation | SiLU |
22
+ | Norm | RMSNorm |
23
+ | Tied Embeddings | Yes |
24
+
25
+ ## Training
26
+
27
+ **Pre-training (4B tokens):**
28
+ - Data: FineWeb-Edu + Cosmopedia (89:11 blend)
29
+ - Schedule: WSD (warmup 2,000 → stable → decay)
30
+ - Peak LR: 7e-4
31
+ - Batch: 16 × 1,024 = 16,384 tokens/step
32
+ - Steps: 244,141
33
+ - Duration: ~102.5 hours (4× RTX 4090)
34
+
35
+ **Post-training (1B tokens):**
36
+ - Data: FineWeb-Edu + Cosmopedia continuation (second 1B slice)
37
+ - Schedule: WSD (warmup 2,000 → stable → decay)
38
+ - Peak LR: 2e-5
39
+ - Batch: 16 × 1,024
40
+ - Steps: 60,975
41
+ - Duration: ~25.9 hours
42
+
43
+ ## Benchmark Results
44
+
45
+ ### Harness (0-shot)
46
+
47
+ | Benchmark | Score |
48
+ |---|---|
49
+ | HellaSwag (acc_norm) | 0.313 |
50
+ | PIQA (acc) | 0.609 |
51
+ | Winogrande (acc) | 0.505 |
52
+ | ARC-Easy (acc_norm) | 0.410 |
53
+ | ARC-Challenge (acc_norm) | 0.272 |
54
+ | OpenBookQA (acc_norm) | 0.290 |
55
+ | BoolQ (acc) | 0.528 |
56
+ | LAMBADA (acc) | 0.195 |
57
+
58
+ ### IFEval (instruction-following)
59
+
60
+ | Model | inst_strict |
61
+ |---|---|
62
+ | **1B post-train** | **0.2338** |
63
+ | v1.1-4B | 0.2266 |
64
+ | v1.1 post-train | 0.2182 |
65
+ | 1B-4B | 0.2170 |
66
+ | v1.1 SFT | 0.2026 |
67
+ | v1.1-8B | 0.1727 |
68
+
69
+ ### SAMSum (Dialogue Summarization)
70
+
71
+ | Model | ROUGE-1 | ROUGE-2 | ROUGE-L |
72
+ |---|---|---|---|
73
+ | **1B post-train (0-shot)** | 9.83 | 0.50 | 7.85 |
74
+ | **1B post-train (fine-tuned, 15ep)** | **28.82** | **8.55** | **24.08** |
75
+ | T5-small (60M) | 35.7 | 13.4 | 31.4 |
76
+
77
+ ## Usage
78
+
79
+ ```python
80
+ from transformers import AutoModelForCausalLM, AutoTokenizer
81
+
82
+ model = AutoModelForCausalLM.from_pretrained(
83
+ "publish_posttrain/",
84
+ trust_remote_code=True,
85
+ torch_dtype="bfloat16",
86
+ device_map="auto",
87
+ )
88
+ tokenizer = AutoTokenizer.from_pretrained("publish_posttrain/", legacy=False)
89
+
90
+ prompt = "The capital of France is"
91
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
92
+ output = model.generate(**inputs, max_new_tokens=20, do_sample=False)
93
+ print(tokenizer.decode(output[0], skip_special_tokens=True))
94
+ ```
95
+
96
+ ## Limitations
97
+
98
+ - **Data budget:** Trained on only 4B tokens (pre-training) + 1B tokens (post-training). Comparable models use 100-1000× more data.
99
+ - **Reasoning:** Limited multi-step reasoning and mathematical capability.
100
+ - **Hallucination:** May generate plausible but incorrect facts.
101
+ - **Context:** Effective context is shorter than the 2,048-token window.
102
+
103
+ ## License
104
+
105
+ MIT License. See [LICENSE](LICENSE) for details.
chat_template.jinja ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {% for message in messages %}
2
+ {% if message['role'] == 'user' %}
3
+ {{ '<|user|>
4
+ ' + message['content'] + eos_token }}
5
+ {% elif message['role'] == 'system' %}
6
+ {{ '<|system|>
7
+ ' + message['content'] + eos_token }}
8
+ {% elif message['role'] == 'assistant' %}
9
+ {{ '<|assistant|>
10
+ ' + message['content'] + eos_token }}
11
+ {% endif %}
12
+ {% if loop.last and add_generation_prompt %}
13
+ {{ '<|assistant|>' }}
14
+ {% endif %}
15
+ {% endfor %}
config.json ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "TinyMixtralForCausalLM"
4
+ ],
5
+ "attention_dropout": 0.0,
6
+ "dtype": "float32",
7
+ "eos_token_id": 2,
8
+ "expert_intermediate_size": 2816,
9
+ "head_dim": 64,
10
+ "hidden_size": 1024,
11
+ "initializer_range": 0.02,
12
+ "max_position_embeddings": 2048,
13
+ "model_type": "tinymixtral",
14
+ "num_attention_heads": 16,
15
+ "num_experts_per_tok": 2,
16
+ "num_hidden_layers": 16,
17
+ "num_key_value_heads": 4,
18
+ "num_local_experts": 8,
19
+ "pad_token_id": 2,
20
+ "rms_norm_eps": 1e-06,
21
+ "rope_theta": 1000000.0,
22
+ "router_aux_loss_coef": 0.01,
23
+ "router_jitter_noise": 0.01,
24
+ "transformers_version": "4.57.3",
25
+ "vocab_size": 32000,
26
+ "auto_map": {
27
+ "AutoConfig": "configuration_tinymixtral.TinyMixtralConfig",
28
+ "AutoModelForCausalLM": "modeling_tinymixtral.TinyMixtralForCausalLM"
29
+ }
30
+ }
configuration_tinymixtral.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (C) Michael Lee (李登淳) 2026. All rights reserved.
2
+ # Open-source under the MIT License. See LICENSE for details.
3
+
4
+ from transformers import PretrainedConfig
5
+
6
+
7
+ class TinyMixtralConfig(PretrainedConfig):
8
+ model_type = "tinymixtral"
9
+
10
+ def __init__(
11
+ self,
12
+ vocab_size: int = 32000,
13
+ hidden_size: int = 896,
14
+ num_hidden_layers: int = 10,
15
+ num_attention_heads: int = 14,
16
+ num_key_value_heads: int = 2,
17
+ head_dim: int = 64,
18
+ max_position_embeddings: int = 2048,
19
+ num_local_experts: int = 6,
20
+ num_experts_per_tok: int = 2,
21
+ expert_intermediate_size: int = 2389,
22
+ router_aux_loss_coef: float = 0.01,
23
+ router_jitter_noise: float = 0.01,
24
+ rms_norm_eps: float = 1e-6,
25
+ rope_theta: float = 1_000_000.0,
26
+ attention_dropout: float = 0.0,
27
+ tie_word_embeddings: bool = True,
28
+ initializer_range: float = 0.02,
29
+ **kwargs,
30
+ ):
31
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
32
+ self.vocab_size = vocab_size
33
+ self.hidden_size = hidden_size
34
+ self.num_hidden_layers = num_hidden_layers
35
+ self.num_attention_heads = num_attention_heads
36
+ self.num_key_value_heads = num_key_value_heads
37
+ self.head_dim = head_dim
38
+ self.max_position_embeddings = max_position_embeddings
39
+ self.num_local_experts = num_local_experts
40
+ self.num_experts_per_tok = num_experts_per_tok
41
+ self.expert_intermediate_size = expert_intermediate_size
42
+ self.router_aux_loss_coef = router_aux_loss_coef
43
+ self.router_jitter_noise = router_jitter_noise
44
+ self.rms_norm_eps = rms_norm_eps
45
+ self.rope_theta = rope_theta
46
+ self.attention_dropout = attention_dropout
47
+ self.initializer_range = initializer_range
generation_config.json ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "max_new_tokens": 256,
3
+ "do_sample": true,
4
+ "temperature": 0.7,
5
+ "top_p": 0.9,
6
+ "eos_token_id": 2,
7
+ "pad_token_id": 2,
8
+ "transformers_version": "4.57.3"
9
+ }
modeling_tinymixtral.py ADDED
@@ -0,0 +1,306 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (C) Michael Lee (李登淳) 2026. All rights reserved.
2
+ # Open-source under the MIT License. See LICENSE for details.
3
+
4
+ from dataclasses import dataclass
5
+ from typing import Optional
6
+
7
+ import torch
8
+ import torch.nn as nn
9
+ import torch.nn.functional as F
10
+ from torch.utils.checkpoint import checkpoint
11
+ from transformers import PreTrainedModel, GenerationMixin
12
+ from transformers.modeling_outputs import ModelOutput
13
+
14
+ from .configuration_tinymixtral import TinyMixtralConfig
15
+
16
+
17
+ # ============================================================
18
+ # Layers
19
+ # ============================================================
20
+
21
+ class RMSNorm(nn.Module):
22
+ def __init__(self, dim: int, eps: float = 1e-6):
23
+ super().__init__()
24
+ self.weight = nn.Parameter(torch.ones(dim))
25
+ self.eps = eps
26
+
27
+ def forward(self, x):
28
+ dtype = x.dtype
29
+ x = x.float()
30
+ norm = x.pow(2).mean(-1, keepdim=True)
31
+ x = x * torch.rsqrt(norm + self.eps)
32
+ return (x * self.weight).to(dtype)
33
+
34
+
35
+ class RotaryEmbedding(nn.Module):
36
+ def __init__(self, dim, max_position_embeddings=2048, theta=10000.0):
37
+ super().__init__()
38
+ self.dim = dim
39
+ self.max_position_embeddings = max_position_embeddings
40
+ self.theta = theta
41
+ self._build_cache()
42
+
43
+ def _build_cache(self):
44
+ inv_freq = 1.0 / (self.theta ** (torch.arange(0, self.dim, 2).float() / self.dim))
45
+ t = torch.arange(self.max_position_embeddings).float()
46
+ freqs = torch.outer(t, inv_freq)
47
+ emb = torch.cat((freqs, freqs), dim=-1)
48
+ self.register_buffer("cos_cached", emb.cos(), persistent=False)
49
+ self.register_buffer("sin_cached", emb.sin(), persistent=False)
50
+
51
+ def forward(self, x, position_ids):
52
+ cos = self.cos_cached[position_ids].unsqueeze(1)
53
+ sin = self.sin_cached[position_ids].unsqueeze(1)
54
+ x_rot = x.float()
55
+ x1, x2 = x_rot.chunk(2, dim=-1)
56
+ rotated = torch.cat((-x2, x1), dim=-1)
57
+ return (x_rot * cos + rotated * sin).to(x.dtype)
58
+
59
+
60
+ class GQAAttention(nn.Module):
61
+ def __init__(self, config):
62
+ super().__init__()
63
+ self.hidden_size = config.hidden_size
64
+ self.num_heads = config.num_attention_heads
65
+ self.num_kv_heads = config.num_key_value_heads
66
+ self.head_dim = config.head_dim
67
+ self.num_groups = self.num_heads // self.num_kv_heads
68
+ assert self.num_heads % self.num_kv_heads == 0
69
+
70
+ self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
71
+ self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
72
+ self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
73
+ self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
74
+ self.rotary_emb = RotaryEmbedding(self.head_dim, config.max_position_embeddings, config.rope_theta)
75
+ self.attention_dropout = config.attention_dropout
76
+
77
+ def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
78
+ B, S, _ = hidden_states.shape
79
+ q = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
80
+ k = self.k_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
81
+ v = self.v_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
82
+
83
+ cache_len = past_key_value[0].shape[2] if past_key_value is not None else 0
84
+ if position_ids is None:
85
+ position_ids = torch.arange(cache_len, cache_len + S, device=hidden_states.device).unsqueeze(0).expand(B, -1)
86
+ q, k = self.rotary_emb(q, position_ids), self.rotary_emb(k, position_ids)
87
+
88
+ if past_key_value is not None:
89
+ k = torch.cat([past_key_value[0], k], dim=2)
90
+ v = torch.cat([past_key_value[1], v], dim=2)
91
+ cache = (k, v) if use_cache else None
92
+ total_len = cache_len + S
93
+
94
+ if attention_mask is not None or cache_len > 0:
95
+ k_exp = k.unsqueeze(2).expand(-1, -1, self.num_groups, -1, -1).reshape(B, self.num_heads, total_len, self.head_dim)
96
+ v_exp = v.unsqueeze(2).expand(-1, -1, self.num_groups, -1, -1).reshape(B, self.num_heads, total_len, self.head_dim)
97
+ causal = torch.tril(torch.ones(S, total_len, device=hidden_states.device, dtype=torch.bool), diagonal=cache_len)
98
+ if attention_mask is not None:
99
+ mask = causal[None, None, :, :] & attention_mask[:, None, None, :]
100
+ else:
101
+ mask = causal[None, None, :, :]
102
+ attn = F.scaled_dot_product_attention(
103
+ q, k_exp, v_exp, attn_mask=mask,
104
+ dropout_p=self.attention_dropout if self.training else 0.0,
105
+ is_causal=False,
106
+ )
107
+ else:
108
+ attn = F.scaled_dot_product_attention(
109
+ q, k, v, attn_mask=None,
110
+ dropout_p=self.attention_dropout if self.training else 0.0,
111
+ is_causal=True,
112
+ enable_gqa=True,
113
+ )
114
+ return self.o_proj(attn.transpose(1, 2).reshape(B, S, -1)), cache
115
+
116
+
117
+ class SparseMoE(nn.Module):
118
+ def __init__(self, config):
119
+ super().__init__()
120
+ self.hidden_size = config.hidden_size
121
+ self.num_experts = config.num_local_experts
122
+ self.top_k = config.num_experts_per_tok
123
+ self.expert_intermediate = config.expert_intermediate_size
124
+ self.jitter_noise = config.router_jitter_noise
125
+ self.aux_loss_coef = config.router_aux_loss_coef
126
+ self.router = nn.Linear(self.hidden_size, self.num_experts, bias=False)
127
+ self.gate_proj = nn.Parameter(torch.empty(self.num_experts, self.expert_intermediate, self.hidden_size))
128
+ self.up_proj = nn.Parameter(torch.empty(self.num_experts, self.expert_intermediate, self.hidden_size))
129
+ self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, self.expert_intermediate))
130
+ self._init_weights()
131
+
132
+ def _init_weights(self, std=0.02):
133
+ nn.init.normal_(self.gate_proj, std=std)
134
+ nn.init.normal_(self.up_proj, std=std)
135
+ nn.init.normal_(self.down_proj, std=std)
136
+
137
+ def forward(self, x):
138
+ B, S, D = x.shape
139
+ x_flat = x.view(-1, D)
140
+ N = B * S
141
+ logits = self.router(x_flat)
142
+ if self.training and self.jitter_noise > 0:
143
+ logits = logits * (1 + torch.randn_like(logits) * self.jitter_noise)
144
+ weights = F.softmax(logits.float(), dim=-1).to(x.dtype)
145
+ w_topk, experts = torch.topk(weights, self.top_k, dim=-1)
146
+ w_topk = w_topk / w_topk.sum(dim=-1, keepdim=True)
147
+
148
+ aux = torch.tensor(0.0, device=x.device, dtype=x.dtype)
149
+ if self.training and self.aux_loss_coef > 0:
150
+ with torch.no_grad():
151
+ mask = F.one_hot(experts, num_classes=self.num_experts).float()
152
+ f_i = mask.mean(dim=(0, 1))
153
+ P_i = weights.mean(dim=0)
154
+ aux = (f_i.detach() * P_i).sum() * self.num_experts
155
+
156
+ flat_experts = experts.view(-1)
157
+ flat_weights = w_topk.view(-1)
158
+ flat_token_idx = torch.arange(N, device=x.device).unsqueeze(1).expand(-1, self.top_k).reshape(-1)
159
+
160
+ sorted_indices = flat_experts.argsort(stable=True)
161
+ sorted_token_idx = flat_token_idx[sorted_indices]
162
+ sorted_weights = flat_weights[sorted_indices]
163
+ sorted_experts = flat_experts[sorted_indices]
164
+
165
+ expert_counts = torch.bincount(sorted_experts, minlength=self.num_experts).tolist()
166
+
167
+ out = torch.zeros(N, D, device=x.device, dtype=x.dtype)
168
+ start = 0
169
+ for e in range(self.num_experts):
170
+ count = expert_counts[e]
171
+ if count == 0:
172
+ continue
173
+ end = start + count
174
+ idx = sorted_token_idx[start:end]
175
+ w = sorted_weights[start:end]
176
+ ts = x_flat[idx]
177
+ gate = F.silu(ts @ self.gate_proj[e].T)
178
+ up = ts @ self.up_proj[e].T
179
+ out.index_add_(0, idx, ((gate * up @ self.down_proj[e].T) * w.unsqueeze(-1)).to(x.dtype))
180
+ start = end
181
+ return out.view(B, S, D), aux
182
+
183
+
184
+ class MoETransformerBlock(nn.Module):
185
+ def __init__(self, config):
186
+ super().__init__()
187
+ self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
188
+ self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
189
+ self.self_attn = GQAAttention(config)
190
+ self.moe = SparseMoE(config)
191
+
192
+ def forward(self, x, attention_mask=None, position_ids=None, past_key_value=None, use_cache=False):
193
+ attn_out, new_cache = self.self_attn(
194
+ self.input_layernorm(x), attention_mask, position_ids, past_key_value, use_cache
195
+ )
196
+ x = x + attn_out
197
+ h, aux = self.moe(self.post_attention_layernorm(x))
198
+ return x + h, aux, new_cache
199
+
200
+
201
+ # ============================================================
202
+ # Causal LM
203
+ # ============================================================
204
+
205
+ @dataclass
206
+ class CausalLMOutputWithPast(ModelOutput):
207
+ loss: Optional[torch.Tensor] = None
208
+ logits: torch.Tensor = None
209
+ past_key_values: Optional[tuple] = None
210
+
211
+
212
+ class TinyMixtralForCausalLM(PreTrainedModel, GenerationMixin):
213
+ config_class = TinyMixtralConfig
214
+ base_model_prefix = "tinymixtral"
215
+ supports_gradient_checkpointing = True
216
+ _no_split_modules = ["MoETransformerBlock"]
217
+ _supports_cache_class = False
218
+ _supports_static_cache = False
219
+
220
+ def _supports_default_dynamic_cache(self):
221
+ return False
222
+
223
+ def __init__(self, config):
224
+ super().__init__(config)
225
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
226
+ self.layers = nn.ModuleList([MoETransformerBlock(config) for _ in range(config.num_hidden_layers)])
227
+ self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
228
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
229
+ if config.tie_word_embeddings:
230
+ self.lm_head.weight = self.embed_tokens.weight
231
+ self._use_activation_checkpointing = False
232
+ self.post_init()
233
+ if getattr(self.config, "eos_token_id", None) is None:
234
+ self.config.eos_token_id = 2
235
+ self.config.pad_token_id = 2
236
+
237
+ def prepare_inputs_for_generation(self, input_ids, attention_mask=None, past_key_values=None, **kwargs):
238
+ return {
239
+ "input_ids": input_ids,
240
+ "attention_mask": attention_mask,
241
+ "past_key_values": past_key_values,
242
+ "use_cache": kwargs.get("use_cache", True),
243
+ }
244
+
245
+ def _reorder_cache(self, past_key_values, beam_idx):
246
+ return tuple(
247
+ tuple(past.index_select(0, beam_idx) for past in layer_past)
248
+ for layer_past in past_key_values
249
+ )
250
+
251
+ def _init_weights(self, module):
252
+ std = self.config.initializer_range
253
+ if isinstance(module, nn.Linear):
254
+ module.weight.data.normal_(mean=0.0, std=std)
255
+ if module.bias is not None:
256
+ module.bias.data.zero_()
257
+ elif isinstance(module, nn.Embedding):
258
+ module.weight.data.normal_(mean=0.0, std=std)
259
+
260
+ def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
261
+ self._use_activation_checkpointing = True
262
+
263
+ def gradient_checkpointing_disable(self):
264
+ self._use_activation_checkpointing = False
265
+
266
+ def forward(self, input_ids, attention_mask=None, labels=None, return_dict=True, past_key_values=None, use_cache=False, **kwargs):
267
+ B, S = input_ids.shape
268
+ past_len = past_key_values[0][0].shape[2] if past_key_values is not None else 0
269
+ if past_len > 0 and S > past_len:
270
+ input_ids = input_ids[:, past_len:]
271
+ S = input_ids.shape[1]
272
+ pos = torch.arange(past_len, past_len + S, device=input_ids.device).unsqueeze(0).expand(B, -1)
273
+
274
+ total_len = past_len + S
275
+ if attention_mask is not None and attention_mask.shape[1] < total_len:
276
+ pad = torch.ones(B, total_len - attention_mask.shape[1],
277
+ dtype=attention_mask.dtype, device=attention_mask.device)
278
+ attention_mask = torch.cat([pad, attention_mask], dim=1)
279
+ cmask = attention_mask.bool() if attention_mask is not None else None
280
+
281
+ h = self.embed_tokens(input_ids)
282
+ total_aux = torch.tensor(0.0, device=input_ids.device, dtype=torch.float32)
283
+ new_caches = []
284
+ for i, layer in enumerate(self.layers):
285
+ layer_cache = past_key_values[i] if past_key_values is not None else None
286
+ if self._use_activation_checkpointing and self.training:
287
+ h, aux, _ = checkpoint(layer, h, cmask, pos, None, False, use_reentrant=False)
288
+ else:
289
+ h, aux, layer_new_cache = layer(h, cmask, pos, layer_cache, use_cache)
290
+ new_caches.append(layer_new_cache)
291
+ total_aux = total_aux + aux
292
+ logits = self.lm_head(self.norm(h)).float()
293
+
294
+ loss = None
295
+ if labels is not None:
296
+ loss = F.cross_entropy(
297
+ logits.reshape(-1, logits.size(-1)),
298
+ labels.reshape(-1),
299
+ ignore_index=-100,
300
+ )
301
+ loss = loss + self.config.router_aux_loss_coef * (total_aux / len(self.layers))
302
+
303
+ past_key_values_out = tuple(new_caches) if use_cache else None
304
+ if not return_dict:
305
+ return (loss, logits, past_key_values_out) if loss is not None else (logits, past_key_values_out)
306
+ return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=past_key_values_out)
pytorch_model.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:55005875e4a997456dd06f4cf416a76f7eb54aca652443df074c34010205530e
3
+ size 4728740711
special_tokens_map.json ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ "pad_token": {
17
+ "content": "</s>",
18
+ "lstrip": false,
19
+ "normalized": false,
20
+ "rstrip": false,
21
+ "single_word": false
22
+ },
23
+ "unk_token": {
24
+ "content": "<unk>",
25
+ "lstrip": false,
26
+ "normalized": false,
27
+ "rstrip": false,
28
+ "single_word": false
29
+ }
30
+ }
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
3
+ size 499723
tokenizer_config.json ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": true,
3
+ "add_eos_token": false,
4
+ "add_prefix_space": null,
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
+ "bos_token": "<s>",
32
+ "clean_up_tokenization_spaces": false,
33
+ "eos_token": "</s>",
34
+ "extra_special_tokens": {},
35
+ "legacy": false,
36
+ "model_max_length": 2048,
37
+ "pad_token": "</s>",
38
+ "padding_side": "right",
39
+ "sp_model_kwargs": {},
40
+ "tokenizer_class": "LlamaTokenizer",
41
+ "unk_token": "<unk>",
42
+ "use_default_system_prompt": false
43
+ }