kgrabko commited on
Commit
1fe3445
·
verified ·
1 Parent(s): d1d15db

Upload 10 files

Browse files
source/JiRack_H12_L12_V50257_D768_MSL8192_FF768x4.py ADDED
@@ -0,0 +1,320 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2025 CMS Manhattan
2
+ # All rights reserved.
3
+ # Author: Konstantin Vladimirovich Grabko
4
+ # Email: grabko@cmsmanhattan.com
5
+ # Phone: +1(516)777-0945
6
+ #
7
+ # This program is free software: you can redistribute it and/or modify
8
+ # it under the terms of the GNU General Public License as published by
9
+ # the Free Software Foundation, version 3 of the License.
10
+ #
11
+ # This program is distributed in the hope that it will be useful,
12
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
13
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
+ # GNU General Public License for more details.
15
+ #
16
+ # You should have received a copy of the GNU General Public License
17
+ # along with this program. If not, see <https://www.gnu.org/licenses/>.
18
+ #
19
+ # Additional terms:
20
+ # Any commercial use or distribution of this software or derivative works
21
+ # requires explicit written permission from the copyright holder.
22
+
23
+ import os
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from typing import Optional, Tuple, List
28
+ import math
29
+ from pathlib import Path
30
+
31
+ # ========================================
32
+ # Model Configuration (GPT-2 Base Style)
33
+ # ========================================
34
+ VOCAB_SIZE = 50257
35
+ MODEL_DIM = 768
36
+ NUM_HEADS = 12
37
+ NUM_LAYERS = 12
38
+ MAX_SEQ_LEN = 8192
39
+ FFN_HIDDEN_DIM = 4 * MODEL_DIM
40
+ HEAD_DIM = MODEL_DIM // NUM_HEADS # 768 / 12 = 64
41
+
42
+ # ИСПРАВЛЕНИЕ: Проверка устройства
43
+ if torch.cuda.is_available():
44
+ device = torch.device("cuda")
45
+ elif hasattr(torch, 'hip') and torch.hip.is_available():
46
+ device = torch.device("cuda")
47
+ else:
48
+ device = torch.device("cpu")
49
+
50
+ # -------------------------------
51
+ # Learned Positional Embedding
52
+ # -------------------------------
53
+ class LearnedPositionalEmbedding(nn.Module):
54
+ def __init__(self, max_seq_len: int, embed_dim: int):
55
+ super().__init__()
56
+ self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
57
+
58
+ def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
59
+ seq_len = x.size(1)
60
+ # УДАЛЕНА Python-проверка для JIT-совместимости, если требуется
61
+ pos = self.pos_emb[pos_offset : pos_offset + seq_len]
62
+ return x + pos.unsqueeze(0)
63
+
64
+
65
+ # -------------------------------
66
+ # MultiHeadAttention (MHA)
67
+ # -------------------------------
68
+ class MultiHeadAttention(nn.Module):
69
+ def __init__(self):
70
+ super().__init__()
71
+ self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
72
+ self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
73
+ self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
74
+ self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
75
+ self.scale = HEAD_DIM ** -0.5
76
+
77
+ def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
78
+ B, T, D = x.shape
79
+
80
+ q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
81
+ k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
82
+ v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
83
+
84
+ # 1. KV-кеш
85
+ seqlen_k_new = k.size(2)
86
+ pos_offset = 0
87
+ if past_kv is not None:
88
+ past_k, past_v = past_kv
89
+ k = torch.cat([past_k, k], dim=2)
90
+ v = torch.cat([past_v, v], dim=2)
91
+ pos_offset = past_k.size(2)
92
+
93
+ seqlen_k = k.size(2)
94
+ new_kv = (k, v)
95
+
96
+ # 2. Расчет внимания
97
+ attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
98
+
99
+ # 3. Маскирование
100
+ if T == seqlen_k_new:
101
+ causal_mask_present = torch.tensor(seqlen_k > 0, dtype=torch.bool, device=x.device)
102
+ if causal_mask_present.item():
103
+ mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
104
+ current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
105
+
106
+ mask[:, :pos_offset] = 0.0
107
+ mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
108
+
109
+ attn = attn + mask[None, None, :, :]
110
+
111
+ # 4. Выход
112
+ attn = F.softmax(attn, dim=-1)
113
+ out = torch.matmul(attn, v)
114
+ out = out.transpose(1, 2).contiguous().view(B, T, D)
115
+ out = self.out_proj(out)
116
+
117
+ return out, new_kv
118
+
119
+
120
+ # -------------------------------
121
+ # FeedForward (GELU, GPT-style)
122
+ # -------------------------------
123
+ class FeedForward(nn.Module):
124
+ def __init__(self):
125
+ super().__init__()
126
+ self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
127
+ self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
128
+
129
+ def forward(self, x):
130
+ return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
131
+
132
+
133
+ # -------------------------------
134
+ # Transformer Block (Post-Norm, GPT-style)
135
+ # -------------------------------
136
+ class TransformerBlock(nn.Module):
137
+ def __init__(self):
138
+ super().__init__()
139
+ self.attn = MultiHeadAttention()
140
+ self.ffn = FeedForward()
141
+ self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
142
+ self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
143
+
144
+ def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
145
+ attn_out, new_kv = self.attn(self.norm1(x), past_kv)
146
+ x = x + attn_out
147
+ x = x + self.ffn(self.norm2(x))
148
+ return x, new_kv
149
+
150
+
151
+ # -------------------------------
152
+ # Главная модель GPTPyTorch (L=12, H=12)
153
+ # -------------------------------
154
+ class GPTPyTorch(nn.Module):
155
+ def __init__(self):
156
+ super().__init__()
157
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
158
+ self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
159
+ self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
160
+ self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
161
+ self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
162
+
163
+ signature = "Konstantin V Gbabko . original author © 2025"
164
+ bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
165
+ self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
166
+ self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
167
+
168
+ self.lm_head.weight = self.token_emb.weight
169
+ self.apply(self._init_weights)
170
+
171
+ def _init_weights(self, module):
172
+ if isinstance(module, nn.Linear):
173
+ std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
174
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
175
+ elif isinstance(module, nn.Embedding):
176
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
177
+ elif isinstance(module, nn.LayerNorm):
178
+ nn.init.zeros_(module.bias)
179
+ nn.init.ones_(module.weight)
180
+
181
+ def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
182
+ B, T = input_ids.shape
183
+ x = self.token_emb(input_ids)
184
+
185
+ pos_offset = 0
186
+ if past_kv is not None and past_kv[0] is not None:
187
+ pos_offset = past_kv[0][0].size(2)
188
+
189
+ x = self.pos_emb(x, pos_offset=pos_offset)
190
+
191
+ # 🎯 ИСПРАВЛЕНИЕ: Инициализация кеша (если T > 1 ИЛИ кеш передан)
192
+ if T > 1 or past_kv is not None:
193
+ new_kv_cache = []
194
+ else:
195
+ new_kv_cache = None
196
+
197
+ current_past = past_kv
198
+
199
+ for i, block in enumerate(self.blocks):
200
+ layer_past = current_past[i] if (current_past and i < len(current_past)) else None
201
+ x, layer_kv = block(x, layer_past)
202
+
203
+ if new_kv_cache is not None:
204
+ new_kv_cache.append(layer_kv)
205
+
206
+ x = self.ln_f(x)
207
+ logits = self.lm_head(x)
208
+ return logits, new_kv_cache
209
+
210
+ @torch.no_grad()
211
+ def generate(
212
+ self,
213
+ input_ids: torch.Tensor,
214
+ max_new_tokens: int = 100,
215
+ temperature: float = 0.8,
216
+ top_p: float = 0.95,
217
+ repetition_penalty: float = 1.0,
218
+ do_sample: bool = True,
219
+ eos_token_id: int = 50256
220
+ ) -> torch.Tensor:
221
+ kv_cache = [None] * NUM_LAYERS
222
+ current_ids = input_ids.clone()
223
+
224
+ for step in range(max_new_tokens):
225
+ if step == 0:
226
+ # Первый проход: обрабатываем весь вход
227
+ input_for_model = current_ids
228
+ else:
229
+ # Инкрементальные проходы: обрабатываем только последний сгенерированный токен
230
+ input_for_model = current_ids[:, -1].unsqueeze(-1) # 🚀 ВОССТАНОВЛЕНА ОБОРОРВАННАЯ ЧАСТЬ
231
+
232
+ logits, kv_cache = self(input_for_model, kv_cache)
233
+ next_token_logits = logits[:, -1, :] # Логиты только для последнего токена
234
+
235
+ # Применение температуры
236
+ if temperature > 0:
237
+ next_token_logits = next_token_logits / temperature
238
+
239
+ # Применение Repetition Penalty
240
+ if repetition_penalty != 1.0:
241
+ for i in range(current_ids.shape[0]):
242
+ unique_tokens = torch.unique(current_ids[i]).tolist()
243
+ for token_id in unique_tokens:
244
+ score = next_token_logits[i, token_id]
245
+ if score < 0:
246
+ next_token_logits[i, token_id] = score * repetition_penalty
247
+ else:
248
+ next_token_logits[i, token_id] = score / repetition_penalty
249
+
250
+ # Top-P сэмплирование
251
+ if do_sample and top_p < 1.0:
252
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
253
+ cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
254
+ sorted_indices_to_remove = cumulative_probs > top_p
255
+
256
+ # Сдвигаем назад, чтобы сохранить первый токен, превысивший top_p
257
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
258
+ sorted_indices_to_remove[:, 0] = False
259
+
260
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
261
+ next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))
262
+
263
+ # Сэмплирование или жадный выбор
264
+ if do_sample and temperature > 0:
265
+ probs = torch.softmax(next_token_logits, dim=-1)
266
+ if torch.isnan(probs).any() or torch.isinf(probs).any():
267
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
268
+ else:
269
+ next_token = torch.multinomial(probs, num_samples=1)
270
+ else:
271
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
272
+
273
+ if next_token.item() == eos_token_id:
274
+ break
275
+
276
+ current_ids = torch.cat([current_ids, next_token], dim=1)
277
+
278
+ return current_ids
279
+
280
+
281
+ if __name__ == "__main__":
282
+ os.makedirs("models", exist_ok=True)
283
+
284
+ model = GPTPyTorch().to(device)
285
+ model.eval()
286
+
287
+ total_params = sum(p.numel() for p in model.parameters())
288
+ print(f"Device: {device}")
289
+ print(f"Total parameters: {total_params / 1e6:.2f}M") # ~200.75M
290
+
291
+ # 1. Проверка первого прохода (T=50)
292
+ input_ids_T50 = torch.randint(0, VOCAB_SIZE, (1, 50), device=device)
293
+ with torch.no_grad():
294
+ logits_50, kv_cache_50 = model(input_ids_T50)
295
+
296
+ # КРИТИЧЕСКАЯ ПРОВЕРКА: kv_cache_50 не должен быть None
297
+ assert kv_cache_50 is not None and len(kv_cache_50) == NUM_LAYERS
298
+
299
+ expected_k_shape = (1, NUM_HEADS, 50, HEAD_DIM)
300
+ assert kv_cache_50[0][0].shape == expected_k_shape
301
+ print(f"Initial logits shape: {logits_50.shape}")
302
+ print(f"Initial KV-cache seqlen: {kv_cache_50[0][0].size(2)} (Head Dim: {kv_cache_50[0][0].size(3)})")
303
+
304
+ # 2. Проверка инкрементального прохода (T=1)
305
+ input_ids_T1 = torch.randint(0, VOCAB_SIZE, (1, 1), device=device)
306
+ with torch.no_grad():
307
+ logits_51, kv_cache_51 = model(input_ids_T1, past_kv=kv_cache_50)
308
+
309
+ # Проверка длины кеша: 50 + 1 = 51
310
+ assert kv_cache_51[0][0].size(2) == 51
311
+ print(f"Incremental logits shape: {logits_51.shape}")
312
+ print(f"Incremental KV-cache seqlen: {kv_cache_51[0][0].size(2)}")
313
+
314
+ # 3. Проверка функции generate
315
+ generated = model.generate(input_ids_T50, max_new_tokens=5, temperature=0.8, top_p=0.9)
316
+ print(f"Generated sequence length (50 + 5): {generated.shape[1]}")
317
+
318
+ save_path = "models/JiRack_H12_L12_V50257_D768_MSL8192_FF768x4.pt"
319
+ torch.save(model.state_dict(), save_path)
320
+ print(f"Model successfully saved to {save_path}")
source/JiRack_H12_L18_V50257_D768_MSL8192_FF768x4.py ADDED
@@ -0,0 +1,318 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2025 CMS Manhattan
2
+ # All rights reserved.
3
+ # Author: Konstantin Vladimirovich Grabko
4
+ # Email: grabko@cmsmanhattan.com
5
+ # Phone: +1(516)777-0945
6
+ #
7
+ # This program is free software: you can redistribute it and/or modify
8
+ # it under the terms of the GNU General Public License as published by
9
+ # the Free Software Foundation, version 3 of the License.
10
+ #
11
+ # This program is distributed in the hope that it will be useful,
12
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
13
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
+ # GNU General Public License for more details.
15
+ #
16
+ # You should have received a copy of the GNU General Public License
17
+ # along with this program. If not, see <https://www.gnu.org/licenses/>.
18
+ #
19
+ # Additional terms:
20
+ # Any commercial use or distribution of this software or derivative works
21
+ # requires explicit written permission from the copyright holder.
22
+
23
+ import os
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from typing import Optional, Tuple, List
28
+ import math
29
+
30
+ # ========================================
31
+ # Model Configuration (GPT-2 Medium Style)
32
+ # ========================================
33
+ VOCAB_SIZE = 50257
34
+ MODEL_DIM = 768
35
+ NUM_HEADS = 12
36
+ NUM_LAYERS = 18 # Increased depth
37
+ MAX_SEQ_LEN = 8192
38
+ FFN_HIDDEN_DIM = 4 * MODEL_DIM
39
+ HEAD_DIM = MODEL_DIM // NUM_HEADS
40
+
41
+ # ИСПРАВЛЕНИЕ: ROCm/HIP-совместимая проверка устройства
42
+ if torch.cuda.is_available():
43
+ device = torch.device("cuda")
44
+ elif hasattr(torch, 'hip') and torch.hip.is_available():
45
+ device = torch.device("cuda")
46
+ else:
47
+ device = torch.device("cpu")
48
+
49
+ # -------------------------------
50
+ # Learned Positional Embedding
51
+ # -------------------------------
52
+ class LearnedPositionalEmbedding(nn.Module):
53
+ def __init__(self, max_seq_len: int, embed_dim: int):
54
+ super().__init__()
55
+ self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
56
+
57
+ def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
58
+ seq_len = x.size(1)
59
+ if pos_offset + seq_len > self.pos_emb.size(0):
60
+ # Проверка на выход за пределы MAX_SEQ_LEN
61
+ raise ValueError("Sequence length exceeds MAX_SEQ_LEN defined in position embedding.")
62
+
63
+ pos = self.pos_emb[pos_offset : pos_offset + seq_len]
64
+ return x + pos.unsqueeze(0)
65
+
66
+
67
+ # -------------------------------
68
+ # MultiHeadAttention (MHA)
69
+ # -------------------------------
70
+ class MultiHeadAttention(nn.Module):
71
+ def __init__(self):
72
+ super().__init__()
73
+ self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
74
+ self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
75
+ self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
76
+ self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
77
+ self.scale = HEAD_DIM ** -0.5
78
+
79
+ def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
80
+ B, T, D = x.shape
81
+
82
+ q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
83
+ k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
84
+ v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
85
+
86
+ # 1. KV-кеш и определение смещения
87
+ pos_offset = 0
88
+ seqlen_k_new = k.size(2)
89
+ if past_kv is not None:
90
+ past_k, past_v = past_kv
91
+ k = torch.cat([past_k, k], dim=2)
92
+ v = torch.cat([past_v, v], dim=2)
93
+ pos_offset = past_k.size(2)
94
+
95
+ seqlen_k = k.size(2) # Общая длина K
96
+ new_kv = (k, v)
97
+
98
+ # 2. Расчет внимания
99
+ attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
100
+
101
+ # 3. КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ МАСКИРОВАНИЯ (Causal Mask)
102
+ if T == seqlen_k_new and seqlen_k > 0:
103
+ # Создаем маску T x seqlen_k
104
+ mask = torch.full((T, seqlen_k),
105
+ float("-inf"),
106
+ device=x.device,
107
+ dtype=attn.dtype)
108
+
109
+ # Разрешаем всем новым токенам видеть все старые токены (past_kv)
110
+ mask[:, :pos_offset] = 0.0
111
+
112
+ # Применяем треугольную маску для текущих T токенов
113
+ current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
114
+ mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
115
+
116
+ # Применяем маску к весам внимания
117
+ attn = attn + mask[None, None, :, :]
118
+
119
+ # 4. Выход
120
+ attn = F.softmax(attn, dim=-1)
121
+ out = torch.matmul(attn, v)
122
+ out = out.transpose(1, 2).contiguous().view(B, T, D)
123
+ out = self.out_proj(out)
124
+
125
+ return out, new_kv
126
+
127
+
128
+ # -------------------------------
129
+ # FeedForward (GELU, GPT-style)
130
+ # -------------------------------
131
+ class FeedForward(nn.Module):
132
+ def __init__(self):
133
+ super().__init__()
134
+ self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
135
+ self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
136
+
137
+ def forward(self, x):
138
+ return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
139
+
140
+
141
+ # -------------------------------
142
+ # Transformer Block (Post-Norm, GPT-style)
143
+ # -------------------------------
144
+ class TransformerBlock(nn.Module):
145
+ def __init__(self):
146
+ super().__init__()
147
+ self.attn = MultiHeadAttention()
148
+ self.ffn = FeedForward()
149
+ # ИСПРАВЛЕНИЕ: Добавлен стандартный eps
150
+ self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
151
+ self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
152
+
153
+ def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
154
+ # Post-Normalization (GPT Style): Input -> Attn(Norm(Input)) + Input -> FFN(Norm(Result)) + Result
155
+ attn_out, new_kv = self.attn(self.norm1(x), past_kv)
156
+ x = x + attn_out
157
+ x = x + self.ffn(self.norm2(x))
158
+ return x, new_kv
159
+
160
+
161
+ # -------------------------------
162
+ # Главная модель GPTPyTorch (18 слоев)
163
+ # -------------------------------
164
+ class GPTPyTorch(nn.Module):
165
+ def __init__(self):
166
+ super().__init__()
167
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
168
+ self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
169
+ self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
170
+ self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
171
+ self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
172
+
173
+ signature = "Konstantin V Gbabko . original author © 2025"
174
+ bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
175
+ self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
176
+ self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
177
+
178
+ self.lm_head.weight = self.token_emb.weight
179
+ self.apply(self._init_weights)
180
+
181
+ def _init_weights(self, module):
182
+ if isinstance(module, nn.Linear):
183
+ # ИСПРАВЛЕНИЕ: Инициализация, масштабированная по глубине сети (TFixup/ReZero style)
184
+ std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
185
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
186
+ elif isinstance(module, nn.Embedding):
187
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
188
+ elif isinstance(module, nn.LayerNorm):
189
+ nn.init.zeros_(module.bias)
190
+ nn.init.ones_(module.weight)
191
+
192
+ def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
193
+ B, T = input_ids.shape
194
+ x = self.token_emb(input_ids)
195
+
196
+ pos_offset = 0
197
+ if past_kv is not None and past_kv[0] is not None:
198
+ pos_offset = past_kv[0][0].size(2)
199
+
200
+ x = self.pos_emb(x, pos_offset=pos_offset)
201
+
202
+ # ИСПРАВЛЕНИЕ: Инициализируем новый кеш
203
+ new_kv_cache = [] if past_kv is not None or T > 1 else None
204
+ current_past = past_kv
205
+
206
+ for i, block in enumerate(self.blocks):
207
+ layer_past = current_past[i] if (current_past and i < len(current_past)) else None
208
+ x, layer_kv = block(x, layer_past)
209
+
210
+ if new_kv_cache is not None:
211
+ new_kv_cache.append(layer_kv)
212
+
213
+ x = self.ln_f(x)
214
+ logits = self.lm_head(x)
215
+ return logits, new_kv_cache
216
+
217
+ @torch.no_grad()
218
+ def generate(
219
+ self,
220
+ input_ids: torch.Tensor,
221
+ max_new_tokens: int = 100,
222
+ temperature: float = 0.8,
223
+ top_p: float = 0.95,
224
+ repetition_penalty: float = 1.0,
225
+ do_sample: bool = True,
226
+ eos_token_id: int = 50256
227
+ ) -> torch.Tensor:
228
+ # ИСПРАВЛЕНИЕ: Инициализируем KV cache как список None кортежей
229
+ kv_cache = [None] * NUM_LAYERS
230
+ current_ids = input_ids.clone()
231
+
232
+ for step in range(max_new_tokens):
233
+ if step == 0:
234
+ input_for_model = current_ids
235
+ else:
236
+ input_for_model = current_ids[:, -1].unsqueeze(-1)
237
+
238
+ logits, kv_cache = self(input_for_model, kv_cache)
239
+ next_token_logits = logits[:, -1, :]
240
+
241
+ if temperature > 0:
242
+ next_token_logits = next_token_logits / temperature
243
+
244
+ # Repetition Penalty (логика сохранена)
245
+ if repetition_penalty != 1.0:
246
+ for i in range(current_ids.shape[0]):
247
+ unique_tokens = torch.unique(current_ids[i]).tolist()
248
+ for token_id in unique_tokens:
249
+ score = next_token_logits[i, token_id]
250
+ if score < 0:
251
+ next_token_logits[i, token_id] = score * repetition_penalty
252
+ else:
253
+ next_token_logits[i, token_id] = score / repetition_penalty
254
+
255
+ # Top-P сэмплирование (логика сохранена)
256
+ if do_sample and top_p < 1.0:
257
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
258
+ cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
259
+ sorted_indices_to_remove = cumulative_probs > top_p
260
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
261
+ sorted_indices_to_remove[:, 0] = False
262
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
263
+ next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))
264
+
265
+ # Сэмплирование
266
+ if do_sample and temperature > 0:
267
+ probs = torch.softmax(next_token_logits, dim=-1)
268
+ if torch.isnan(probs).any() or torch.isinf(probs).any():
269
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
270
+ else:
271
+ next_token = torch.multinomial(probs, num_samples=1)
272
+ else:
273
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
274
+
275
+ if next_token.item() == eos_token_id:
276
+ break
277
+
278
+ current_ids = torch.cat([current_ids, next_token], dim=1)
279
+
280
+ return current_ids
281
+
282
+
283
+ if __name__ == "__main__":
284
+ os.makedirs("models", exist_ok=True)
285
+
286
+ model = GPTPyTorch().to(device)
287
+ model.eval()
288
+
289
+ total_params = sum(p.numel() for p in model.parameters())
290
+ print(f"Device: {device}")
291
+ print(f"Total parameters: {total_params / 1e6:.2f}M") # ~167.33M
292
+
293
+ # 1. Проверка первого прохода (T=50)
294
+ input_ids_T50 = torch.randint(0, VOCAB_SIZE, (1, 50), device=device)
295
+ with torch.no_grad():
296
+ logits_50, kv_cache_50 = model(input_ids_T50)
297
+ expected_k_shape = (1, NUM_HEADS, 50, HEAD_DIM)
298
+ assert kv_cache_50[0][0].shape == expected_k_shape
299
+ print(f"Initial logits shape: {logits_50.shape}")
300
+ print(f"Initial KV-cache seqlen: {kv_cache_50[0][0].size(2)}")
301
+
302
+ # 2. Проверка инкрементального прохода (T=1)
303
+ input_ids_T1 = torch.randint(0, VOCAB_SIZE, (1, 1), device=device)
304
+ with torch.no_grad():
305
+ logits_51, kv_cache_51 = model(input_ids_T1, past_kv=kv_cache_50)
306
+
307
+ # Проверка длины кеша: 50 + 1 = 51
308
+ assert kv_cache_51[0][0].size(2) == 51
309
+ print(f"Incremental logits shape: {logits_51.shape}")
310
+ print(f"Incremental KV-cache seqlen: {kv_cache_51[0][0].size(2)}")
311
+
312
+ # 3. Проверка функции generate
313
+ generated = model.generate(input_ids_T50, max_new_tokens=5, temperature=0.8, top_p=0.9)
314
+ print(f"Generated sequence length (50 + 5): {generated.shape[1]}")
315
+
316
+ save_path = "models/JiRack_GPT_L18_PostNorm_fixed.pt"
317
+ torch.save(model.state_dict(), save_path)
318
+ print(f"Model successfully saved to {save_path}")
source/JiRack_H12_L24_V50257_D768_MSL8192_FF768x4.py ADDED
@@ -0,0 +1,317 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2025 CMS Manhattan
2
+ # All rights reserved.
3
+ # Author: Konstantin Vladimirovich Grabko
4
+ # Email: grabko@cmsmanhattan.com
5
+ # Phone: +1(516)777-0945
6
+ #
7
+ # This program is free software: you can redistribute it and/or modify
8
+ # it under the terms of the GNU General Public License as published by
9
+ # the Free Software Foundation, version 3 of the License.
10
+ #
11
+ # This program is distributed in the hope that it will be useful,
12
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
13
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
+ # GNU General Public License for more details.
15
+ #
16
+ # You should have received a copy of the GNU General Public License
17
+ # along with this program. If not, see <https://www.gnu.org/licenses/>.
18
+ #
19
+ # Additional terms:
20
+ # Any commercial use or distribution of this software or derivative works
21
+ # requires explicit written permission from the copyright holder.
22
+
23
+ import os
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from typing import Optional, Tuple, List
28
+ import math
29
+
30
+ # ========================================
31
+ # Model Configuration (GPT-2 Large Style)
32
+ # ========================================
33
+ VOCAB_SIZE = 50257
34
+ MODEL_DIM = 768
35
+ NUM_HEADS = 12
36
+ NUM_LAYERS = 24 # Increased depth (GPT-2 Large equivalent)
37
+ MAX_SEQ_LEN = 8192
38
+ FFN_HIDDEN_DIM = 4 * MODEL_DIM
39
+ HEAD_DIM = MODEL_DIM // NUM_HEADS
40
+
41
+ # ИСПРАВЛЕНИЕ: ROCm/HIP-совместимая проверка устройства
42
+ if torch.cuda.is_available():
43
+ device = torch.device("cuda")
44
+ elif hasattr(torch, 'hip') and torch.hip.is_available():
45
+ device = torch.device("cuda")
46
+ else:
47
+ device = torch.device("cpu")
48
+
49
+ # -------------------------------
50
+ # Learned Positional Embedding
51
+ # -------------------------------
52
+ class LearnedPositionalEmbedding(nn.Module):
53
+ def __init__(self, max_seq_len: int, embed_dim: int):
54
+ super().__init__()
55
+ self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
56
+
57
+ def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
58
+ seq_len = x.size(1)
59
+ if pos_offset + seq_len > self.pos_emb.size(0):
60
+ raise ValueError("Sequence length exceeds MAX_SEQ_LEN defined in position embedding.")
61
+
62
+ pos = self.pos_emb[pos_offset : pos_offset + seq_len]
63
+ return x + pos.unsqueeze(0)
64
+
65
+
66
+ # -------------------------------
67
+ # MultiHeadAttention (MHA)
68
+ # -------------------------------
69
+ class MultiHeadAttention(nn.Module):
70
+ def __init__(self):
71
+ super().__init__()
72
+ self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
73
+ self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
74
+ self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
75
+ self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
76
+ self.scale = HEAD_DIM ** -0.5
77
+
78
+ def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
79
+ B, T, D = x.shape
80
+
81
+ q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
82
+ k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
83
+ v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
84
+
85
+ # 1. KV-кеш и определение смещения
86
+ pos_offset = 0
87
+ seqlen_k_new = k.size(2)
88
+ if past_kv is not None:
89
+ past_k, past_v = past_kv
90
+ k = torch.cat([past_k, k], dim=2)
91
+ v = torch.cat([past_v, v], dim=2)
92
+ pos_offset = past_k.size(2)
93
+
94
+ seqlen_k = k.size(2) # Общая длина K
95
+ new_kv = (k, v)
96
+
97
+ # 2. Расчет внимания
98
+ attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
99
+
100
+ # 3. КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ МАСКИРОВАНИЯ (Causal Mask)
101
+ if T == seqlen_k_new and seqlen_k > 0:
102
+ # Создаем маску T x seqlen_k
103
+ mask = torch.full((T, seqlen_k),
104
+ float("-inf"),
105
+ device=x.device,
106
+ dtype=attn.dtype)
107
+
108
+ # Разрешаем всем новым токенам видеть все старые токены (past_kv)
109
+ mask[:, :pos_offset] = 0.0
110
+
111
+ # Применяем треугольную маску для текущих T токенов
112
+ current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
113
+ mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
114
+
115
+ # Применяем маску к весам внимания
116
+ attn = attn + mask[None, None, :, :]
117
+
118
+ # 4. Выход
119
+ attn = F.softmax(attn, dim=-1)
120
+ out = torch.matmul(attn, v)
121
+ out = out.transpose(1, 2).contiguous().view(B, T, D)
122
+ out = self.out_proj(out)
123
+
124
+ return out, new_kv
125
+
126
+
127
+ # -------------------------------
128
+ # FeedForward (GELU, GPT-style)
129
+ # -------------------------------
130
+ class FeedForward(nn.Module):
131
+ def __init__(self):
132
+ super().__init__()
133
+ self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
134
+ self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
135
+
136
+ def forward(self, x):
137
+ return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
138
+
139
+
140
+ # -------------------------------
141
+ # Transformer Block (Post-Norm, GPT-style)
142
+ # -------------------------------
143
+ class TransformerBlock(nn.Module):
144
+ def __init__(self):
145
+ super().__init__()
146
+ self.attn = MultiHeadAttention()
147
+ self.ffn = FeedForward()
148
+ # ИСПРАВЛЕНИЕ: Добавлен стандартный eps
149
+ self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
150
+ self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
151
+
152
+ def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
153
+ # Post-Normalization (GPT Style)
154
+ attn_out, new_kv = self.attn(self.norm1(x), past_kv)
155
+ x = x + attn_out
156
+ x = x + self.ffn(self.norm2(x))
157
+ return x, new_kv
158
+
159
+
160
+ # -------------------------------
161
+ # Главная модель GPTPyTorch (24 слоя)
162
+ # -------------------------------
163
+ class GPTPyTorch(nn.Module):
164
+ def __init__(self):
165
+ super().__init__()
166
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
167
+ self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
168
+ self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
169
+ self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
170
+ self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
171
+
172
+ signature = "Konstantin V Gbabko . original author © 2025"
173
+ bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
174
+ self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
175
+ self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
176
+
177
+ self.lm_head.weight = self.token_emb.weight
178
+ self.apply(self._init_weights)
179
+
180
+ def _init_weights(self, module):
181
+ if isinstance(module, nn.Linear):
182
+ # ИСПРАВЛЕНИЕ: Инициализация, масштабированная по глубине сети (TFixup/ReZero style)
183
+ # Стандартное отклонение уменьшается с ростом числа слоев
184
+ std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
185
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
186
+ elif isinstance(module, nn.Embedding):
187
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
188
+ elif isinstance(module, nn.LayerNorm):
189
+ nn.init.zeros_(module.bias)
190
+ nn.init.ones_(module.weight)
191
+
192
+ def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
193
+ B, T = input_ids.shape
194
+ x = self.token_emb(input_ids)
195
+
196
+ pos_offset = 0
197
+ if past_kv is not None and past_kv[0] is not None:
198
+ pos_offset = past_kv[0][0].size(2)
199
+
200
+ x = self.pos_emb(x, pos_offset=pos_offset)
201
+
202
+ # Инициализация нового кеша
203
+ new_kv_cache = [] if past_kv is not None or T > 1 else None
204
+ current_past = past_kv
205
+
206
+ for i, block in enumerate(self.blocks):
207
+ layer_past = current_past[i] if (current_past and i < len(current_past)) else None
208
+ x, layer_kv = block(x, layer_past)
209
+
210
+ if new_kv_cache is not None:
211
+ new_kv_cache.append(layer_kv)
212
+
213
+ x = self.ln_f(x)
214
+ logits = self.lm_head(x)
215
+ return logits, new_kv_cache
216
+
217
+ @torch.no_grad()
218
+ def generate(
219
+ self,
220
+ input_ids: torch.Tensor,
221
+ max_new_tokens: int = 100,
222
+ temperature: float = 0.8,
223
+ top_p: float = 0.95,
224
+ repetition_penalty: float = 1.0,
225
+ do_sample: bool = True,
226
+ eos_token_id: int = 50256
227
+ ) -> torch.Tensor:
228
+ kv_cache = [None] * NUM_LAYERS
229
+ current_ids = input_ids.clone()
230
+
231
+ for step in range(max_new_tokens):
232
+ if step == 0:
233
+ input_for_model = current_ids
234
+ else:
235
+ input_for_model = current_ids[:, -1].unsqueeze(-1)
236
+
237
+ logits, kv_cache = self(input_for_model, kv_cache)
238
+ next_token_logits = logits[:, -1, :]
239
+
240
+ if temperature > 0:
241
+ next_token_logits = next_token_logits / temperature
242
+
243
+ # Repetition Penalty (логика сохранена)
244
+ if repetition_penalty != 1.0:
245
+ for i in range(current_ids.shape[0]):
246
+ unique_tokens = torch.unique(current_ids[i]).tolist()
247
+ for token_id in unique_tokens:
248
+ score = next_token_logits[i, token_id]
249
+ if score < 0:
250
+ next_token_logits[i, token_id] = score * repetition_penalty
251
+ else:
252
+ next_token_logits[i, token_id] = score / repetition_penalty
253
+
254
+ # Top-P сэмплирование (логика сохранена)
255
+ if do_sample and top_p < 1.0:
256
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
257
+ cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
258
+ sorted_indices_to_remove = cumulative_probs > top_p
259
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
260
+ sorted_indices_to_remove[:, 0] = False
261
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
262
+ next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))
263
+
264
+ # Сэмплирование
265
+ if do_sample and temperature > 0:
266
+ probs = torch.softmax(next_token_logits, dim=-1)
267
+ if torch.isnan(probs).any() or torch.isinf(probs).any():
268
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
269
+ else:
270
+ next_token = torch.multinomial(probs, num_samples=1)
271
+ else:
272
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
273
+
274
+ if next_token.item() == eos_token_id:
275
+ break
276
+
277
+ current_ids = torch.cat([current_ids, next_token], dim=1)
278
+
279
+ return current_ids
280
+
281
+
282
+ if __name__ == "__main__":
283
+ os.makedirs("models", exist_ok=True)
284
+
285
+ model = GPTPyTorch().to(device)
286
+ model.eval()
287
+
288
+ total_params = sum(p.numel() for p in model.parameters())
289
+ print(f"Device: {device}")
290
+ print(f"Total parameters: {total_params / 1e6:.2f}M") # ~221.75M
291
+
292
+ # 1. Проверка первого прохода (T=50)
293
+ input_ids_T50 = torch.randint(0, VOCAB_SIZE, (1, 50), device=device)
294
+ with torch.no_grad():
295
+ logits_50, kv_cache_50 = model(input_ids_T50)
296
+ expected_k_shape = (1, NUM_HEADS, 50, HEAD_DIM)
297
+ assert kv_cache_50[0][0].shape == expected_k_shape
298
+ print(f"Initial logits shape: {logits_50.shape}")
299
+ print(f"Initial KV-cache seqlen: {kv_cache_50[0][0].size(2)}")
300
+
301
+ # 2. Проверка инкрементального прохода (T=1)
302
+ input_ids_T1 = torch.randint(0, VOCAB_SIZE, (1, 1), device=device)
303
+ with torch.no_grad():
304
+ logits_51, kv_cache_51 = model(input_ids_T1, past_kv=kv_cache_50)
305
+
306
+ # Проверка длины кеша: 50 + 1 = 51
307
+ assert kv_cache_51[0][0].size(2) == 51
308
+ print(f"Incremental logits shape: {logits_51.shape}")
309
+ print(f"Incremental KV-cache seqlen: {kv_cache_51[0][0].size(2)}")
310
+
311
+ # 3. Проверка функции generate
312
+ generated = model.generate(input_ids_T50, max_new_tokens=5, temperature=0.8, top_p=0.9)
313
+ print(f"Generated sequence length (50 + 5): {generated.shape[1]}")
314
+
315
+ save_path = "models/JiRack_GPT_L24_PostNorm_fixed.pt"
316
+ torch.save(model.state_dict(), save_path)
317
+ print(f"Model successfully saved to {save_path}")
source/JiRack_H12_L32_V50257_D768_MSL8192_FF768x4.py ADDED
@@ -0,0 +1,317 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2025 CMS Manhattan
2
+ # All rights reserved.
3
+ # Author: Konstantin Vladimirovich Grabko
4
+ # Email: grabko@cmsmanhattan.com
5
+ # Phone: +1(516)777-0945
6
+ #
7
+ # This program is free software: you can redistribute it and/or modify
8
+ # it under the terms of the GNU General Public License as published by
9
+ # the Free Software Foundation, version 3 of the License.
10
+ #
11
+ # This program is distributed in the hope that it will be useful,
12
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
13
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
+ # GNU General Public License for more details.
15
+ #
16
+ # You should have received a copy of the GNU General Public License
17
+ # along with this program. If not, see <https://www.gnu.org/licenses/>.
18
+ #
19
+ # Additional terms:
20
+ # Any commercial use or distribution of this software or derivative works
21
+ # requires explicit written permission from the copyright holder.
22
+
23
+ import os
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from typing import Optional, Tuple, List
28
+ import math
29
+
30
+ # ========================================
31
+ # Model Configuration (GPT-2 XL Style)
32
+ # ========================================
33
+ VOCAB_SIZE = 50257
34
+ MODEL_DIM = 768
35
+ NUM_HEADS = 12
36
+ NUM_LAYERS = 32 # Increased depth (GPT-2 XL equivalent)
37
+ MAX_SEQ_LEN = 8192
38
+ FFN_HIDDEN_DIM = 4 * MODEL_DIM
39
+ HEAD_DIM = MODEL_DIM // NUM_HEADS
40
+
41
+ # ИСПРАВЛЕНИЕ: ROCm/HIP-совместимая проверка устройства
42
+ if torch.cuda.is_available():
43
+ device = torch.device("cuda")
44
+ elif hasattr(torch, 'hip') and torch.hip.is_available():
45
+ device = torch.device("cuda")
46
+ else:
47
+ device = torch.device("cpu")
48
+
49
+ # -------------------------------
50
+ # Learned Positional Embedding
51
+ # -------------------------------
52
+ class LearnedPositionalEmbedding(nn.Module):
53
+ def __init__(self, max_seq_len: int, embed_dim: int):
54
+ super().__init__()
55
+ self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
56
+
57
+ def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
58
+ seq_len = x.size(1)
59
+ if pos_offset + seq_len > self.pos_emb.size(0):
60
+ raise ValueError("Sequence length exceeds MAX_SEQ_LEN defined in position embedding.")
61
+
62
+ pos = self.pos_emb[pos_offset : pos_offset + seq_len]
63
+ return x + pos.unsqueeze(0)
64
+
65
+
66
+ # -------------------------------
67
+ # MultiHeadAttention (MHA)
68
+ # -------------------------------
69
+ class MultiHeadAttention(nn.Module):
70
+ def __init__(self):
71
+ super().__init__()
72
+ self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
73
+ self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
74
+ self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
75
+ self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
76
+ self.scale = HEAD_DIM ** -0.5
77
+
78
+ def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
79
+ B, T, D = x.shape
80
+
81
+ q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
82
+ k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
83
+ v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
84
+
85
+ # 1. KV-кеш и определение смещения
86
+ pos_offset = 0
87
+ seqlen_k_new = k.size(2)
88
+ if past_kv is not None:
89
+ past_k, past_v = past_kv
90
+ k = torch.cat([past_k, k], dim=2)
91
+ v = torch.cat([past_v, v], dim=2)
92
+ pos_offset = past_k.size(2)
93
+
94
+ seqlen_k = k.size(2) # Общая длина K
95
+ new_kv = (k, v)
96
+
97
+ # 2. Расчет внимания
98
+ attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
99
+
100
+ # 3. КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ МАСКИРОВАНИЯ (Causal Mask)
101
+ if T == seqlen_k_new and seqlen_k > 0:
102
+ # Создаем маску T x seqlen_k
103
+ mask = torch.full((T, seqlen_k),
104
+ float("-inf"),
105
+ device=x.device,
106
+ dtype=attn.dtype)
107
+
108
+ # Разрешаем всем новым токенам видеть все старые токены (past_kv)
109
+ mask[:, :pos_offset] = 0.0
110
+
111
+ # Применяем треугольную маску для текущих T токенов
112
+ current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
113
+ mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
114
+
115
+ # Применяем маску к весам внимания
116
+ attn = attn + mask[None, None, :, :]
117
+
118
+ # 4. Выход
119
+ attn = F.softmax(attn, dim=-1)
120
+ out = torch.matmul(attn, v)
121
+ out = out.transpose(1, 2).contiguous().view(B, T, D)
122
+ out = self.out_proj(out)
123
+
124
+ return out, new_kv
125
+
126
+
127
+ # -------------------------------
128
+ # FeedForward (GELU, GPT-style)
129
+ # -------------------------------
130
+ class FeedForward(nn.Module):
131
+ def __init__(self):
132
+ super().__init__()
133
+ self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
134
+ self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
135
+
136
+ def forward(self, x):
137
+ return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
138
+
139
+
140
+ # -------------------------------
141
+ # Transformer Block (Post-Norm, GPT-style)
142
+ # -------------------------------
143
+ class TransformerBlock(nn.Module):
144
+ def __init__(self):
145
+ super().__init__()
146
+ self.attn = MultiHeadAttention()
147
+ self.ffn = FeedForward()
148
+ # ИСПРАВЛЕНИЕ: Добавлен стандартный eps
149
+ self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
150
+ self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
151
+
152
+ def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
153
+ # Post-Normalization (GPT Style)
154
+ attn_out, new_kv = self.attn(self.norm1(x), past_kv)
155
+ x = x + attn_out
156
+ x = x + self.ffn(self.norm2(x))
157
+ return x, new_kv
158
+
159
+
160
+ # -------------------------------
161
+ # Главная модель GPTPyTorch (32 слоя)
162
+ # -------------------------------
163
+ class GPTPyTorch(nn.Module):
164
+ def __init__(self):
165
+ super().__init__()
166
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
167
+ self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
168
+ self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
169
+ self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
170
+ self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
171
+
172
+ signature = "Konstantin V Gbabko . original author © 2025"
173
+ bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
174
+ self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
175
+ self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
176
+
177
+ self.lm_head.weight = self.token_emb.weight
178
+ self.apply(self._init_weights)
179
+
180
+ def _init_weights(self, module):
181
+ if isinstance(module, nn.Linear):
182
+ # ИСПРАВЛЕНИЕ: Инициализация, масштабированная по глубине сети (TFixup/ReZero style).
183
+ # Critical for NUM_LAYERS = 32
184
+ std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
185
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
186
+ elif isinstance(module, nn.Embedding):
187
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
188
+ elif isinstance(module, nn.LayerNorm):
189
+ nn.init.zeros_(module.bias)
190
+ nn.init.ones_(module.weight)
191
+
192
+ def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
193
+ B, T = input_ids.shape
194
+ x = self.token_emb(input_ids)
195
+
196
+ pos_offset = 0
197
+ if past_kv is not None and past_kv[0] is not None:
198
+ pos_offset = past_kv[0][0].size(2)
199
+
200
+ x = self.pos_emb(x, pos_offset=pos_offset)
201
+
202
+ # Инициализация нового кеша
203
+ new_kv_cache = [] if past_kv is not None or T > 1 else None
204
+ current_past = past_kv
205
+
206
+ for i, block in enumerate(self.blocks):
207
+ layer_past = current_past[i] if (current_past and i < len(current_past)) else None
208
+ x, layer_kv = block(x, layer_past)
209
+
210
+ if new_kv_cache is not None:
211
+ new_kv_cache.append(layer_kv)
212
+
213
+ x = self.ln_f(x)
214
+ logits = self.lm_head(x)
215
+ return logits, new_kv_cache
216
+
217
+ @torch.no_grad()
218
+ def generate(
219
+ self,
220
+ input_ids: torch.Tensor,
221
+ max_new_tokens: int = 100,
222
+ temperature: float = 0.8,
223
+ top_p: float = 0.95,
224
+ repetition_penalty: float = 1.0,
225
+ do_sample: bool = True,
226
+ eos_token_id: int = 50256
227
+ ) -> torch.Tensor:
228
+ kv_cache = [None] * NUM_LAYERS
229
+ current_ids = input_ids.clone()
230
+
231
+ for step in range(max_new_tokens):
232
+ if step == 0:
233
+ input_for_model = current_ids
234
+ else:
235
+ input_for_model = current_ids[:, -1].unsqueeze(-1)
236
+
237
+ logits, kv_cache = self(input_for_model, kv_cache)
238
+ next_token_logits = logits[:, -1, :]
239
+
240
+ if temperature > 0:
241
+ next_token_logits = next_token_logits / temperature
242
+
243
+ # Repetition Penalty (логика сохранена)
244
+ if repetition_penalty != 1.0:
245
+ for i in range(current_ids.shape[0]):
246
+ unique_tokens = torch.unique(current_ids[i]).tolist()
247
+ for token_id in unique_tokens:
248
+ score = next_token_logits[i, token_id]
249
+ if score < 0:
250
+ next_token_logits[i, token_id] = score * repetition_penalty
251
+ else:
252
+ next_token_logits[i, token_id] = score / repetition_penalty
253
+
254
+ # Top-P сэмплирование (логика сохранена)
255
+ if do_sample and top_p < 1.0:
256
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
257
+ cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
258
+ sorted_indices_to_remove = cumulative_probs > top_p
259
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
260
+ sorted_indices_to_remove[:, 0] = False
261
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
262
+ next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))
263
+
264
+ # Сэмплирование
265
+ if do_sample and temperature > 0:
266
+ probs = torch.softmax(next_token_logits, dim=-1)
267
+ if torch.isnan(probs).any() or torch.isinf(probs).any():
268
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
269
+ else:
270
+ next_token = torch.multinomial(probs, num_samples=1)
271
+ else:
272
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
273
+
274
+ if next_token.item() == eos_token_id:
275
+ break
276
+
277
+ current_ids = torch.cat([current_ids, next_token], dim=1)
278
+
279
+ return current_ids
280
+
281
+
282
+ if __name__ == "__main__":
283
+ os.makedirs("models", exist_ok=True)
284
+
285
+ model = GPTPyTorch().to(device)
286
+ model.eval()
287
+
288
+ total_params = sum(p.numel() for p in model.parameters())
289
+ print(f"Device: {device}")
290
+ print(f"Total parameters: {total_params / 1e6:.2f}M") # ~295.1M
291
+
292
+ # 1. Проверка первого прохода (T=50)
293
+ input_ids_T50 = torch.randint(0, VOCAB_SIZE, (1, 50), device=device)
294
+ with torch.no_grad():
295
+ logits_50, kv_cache_50 = model(input_ids_T50)
296
+ expected_k_shape = (1, NUM_HEADS, 50, HEAD_DIM)
297
+ assert kv_cache_50[0][0].shape == expected_k_shape
298
+ print(f"Initial logits shape: {logits_50.shape}")
299
+ print(f"Initial KV-cache seqlen: {kv_cache_50[0][0].size(2)}")
300
+
301
+ # 2. Проверка инкрементального прохода (T=1)
302
+ input_ids_T1 = torch.randint(0, VOCAB_SIZE, (1, 1), device=device)
303
+ with torch.no_grad():
304
+ logits_51, kv_cache_51 = model(input_ids_T1, past_kv=kv_cache_50)
305
+
306
+ # Проверка длины кеша: 50 + 1 = 51
307
+ assert kv_cache_51[0][0].size(2) == 51
308
+ print(f"Incremental logits shape: {logits_51.shape}")
309
+ print(f"Incremental KV-cache seqlen: {kv_cache_51[0][0].size(2)}")
310
+
311
+ # 3. Проверка функции generate
312
+ generated = model.generate(input_ids_T50, max_new_tokens=5, temperature=0.8, top_p=0.9)
313
+ print(f"Generated sequence length (50 + 5): {generated.shape[1]}")
314
+
315
+ save_path = "models/JiRack_GPT_L32_PostNorm_fixed.pt"
316
+ torch.save(model.state_dict(), save_path)
317
+ print(f"Model successfully saved to {save_path}")
source/JiRack_H12_L6_V50257_D768_MSL8192_FF768x4.py ADDED
@@ -0,0 +1,316 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2025 CMS Manhattan
2
+ # All rights reserved.
3
+ # Author: Konstantin Vladimirovich Grabko
4
+ # Email: grabko@cmsmanhattan.com
5
+ # Phone: +1(516)777-0945
6
+ #
7
+ # This program is free software: you can redistribute it and/or modify
8
+ # it under the terms of the GNU General Public License as published by
9
+ # the Free Software Foundation, version 3 of the License.
10
+ #
11
+ # This program is distributed in the hope that it will be useful,
12
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
13
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
+ # GNU General Public License for more details.
15
+ #
16
+ # You should have received a copy of the GNU General Public License
17
+ # along with this program. If not, see <https://www.gnu.org/licenses/>.
18
+ #
19
+ # Additional terms:
20
+ # Any commercial use or distribution of this software or derivative works
21
+ # requires explicit written permission from the copyright holder.
22
+
23
+ import os
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from typing import Optional, Tuple, List
28
+ import math
29
+
30
+ # ========================================
31
+ # Model Configuration (GPT-2 Small Style)
32
+ # ========================================
33
+ VOCAB_SIZE = 50257
34
+ MODEL_DIM = 768
35
+ NUM_HEADS = 12
36
+ NUM_LAYERS = 6 # Back to 6 layers
37
+ MAX_SEQ_LEN = 8192
38
+ FFN_HIDDEN_DIM = 4 * MODEL_DIM
39
+ HEAD_DIM = MODEL_DIM // NUM_HEADS
40
+
41
+ # ИСПРАВЛЕНИЕ: ROCm/HIP-совместимая проверка устройства
42
+ if torch.cuda.is_available():
43
+ device = torch.device("cuda")
44
+ elif hasattr(torch, 'hip') and torch.hip.is_available():
45
+ device = torch.device("cuda")
46
+ else:
47
+ device = torch.device("cpu")
48
+
49
+ # -------------------------------
50
+ # Learned Positional Embedding
51
+ # -------------------------------
52
+ class LearnedPositionalEmbedding(nn.Module):
53
+ def __init__(self, max_seq_len: int, embed_dim: int):
54
+ super().__init__()
55
+ self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
56
+
57
+ def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
58
+ seq_len = x.size(1)
59
+ if pos_offset + seq_len > self.pos_emb.size(0):
60
+ raise ValueError("Sequence length exceeds MAX_SEQ_LEN defined in position embedding.")
61
+
62
+ pos = self.pos_emb[pos_offset : pos_offset + seq_len]
63
+ return x + pos.unsqueeze(0)
64
+
65
+
66
+ # -------------------------------
67
+ # MultiHeadAttention (MHA)
68
+ # -------------------------------
69
+ class MultiHeadAttention(nn.Module):
70
+ def __init__(self):
71
+ super().__init__()
72
+ self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
73
+ self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
74
+ self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
75
+ self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
76
+ self.scale = HEAD_DIM ** -0.5
77
+
78
+ def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
79
+ B, T, D = x.shape
80
+
81
+ q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
82
+ k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
83
+ v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
84
+
85
+ # 1. KV-кеш и определение смещения
86
+ pos_offset = 0
87
+ seqlen_k_new = k.size(2)
88
+ if past_kv is not None:
89
+ past_k, past_v = past_kv
90
+ k = torch.cat([past_k, k], dim=2)
91
+ v = torch.cat([past_v, v], dim=2)
92
+ pos_offset = past_k.size(2)
93
+
94
+ seqlen_k = k.size(2) # Общая длина K
95
+ new_kv = (k, v)
96
+
97
+ # 2. Расчет внимания
98
+ attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
99
+
100
+ # 3. КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ МАСКИРОВАНИЯ (Causal Mask)
101
+ if T == seqlen_k_new and seqlen_k > 0:
102
+ # Создаем маску T x seqlen_k
103
+ mask = torch.full((T, seqlen_k),
104
+ float("-inf"),
105
+ device=x.device,
106
+ dtype=attn.dtype)
107
+
108
+ # Разрешаем всем новым токенам видеть все старые токены (past_kv)
109
+ mask[:, :pos_offset] = 0.0
110
+
111
+ # Применяем треугольную маску для текущих T токенов
112
+ current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
113
+ mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
114
+
115
+ # Применяем маску к весам внимания
116
+ attn = attn + mask[None, None, :, :]
117
+
118
+ # 4. Выход
119
+ attn = F.softmax(attn, dim=-1)
120
+ out = torch.matmul(attn, v)
121
+ out = out.transpose(1, 2).contiguous().view(B, T, D)
122
+ out = self.out_proj(out)
123
+
124
+ return out, new_kv
125
+
126
+
127
+ # -------------------------------
128
+ # FeedForward (GELU, GPT-style)
129
+ # -------------------------------
130
+ class FeedForward(nn.Module):
131
+ def __init__(self):
132
+ super().__init__()
133
+ self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
134
+ self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
135
+
136
+ def forward(self, x):
137
+ return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
138
+
139
+
140
+ # -------------------------------
141
+ # Transformer Block (Post-Norm, GPT-style)
142
+ # -------------------------------
143
+ class TransformerBlock(nn.Module):
144
+ def __init__(self):
145
+ super().__init__()
146
+ self.attn = MultiHeadAttention()
147
+ self.ffn = FeedForward()
148
+ # ИСПРАВЛЕНИЕ: Добавлен стандартный eps
149
+ self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
150
+ self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
151
+
152
+ def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
153
+ # Post-Normalization (GPT Style)
154
+ attn_out, new_kv = self.attn(self.norm1(x), past_kv)
155
+ x = x + attn_out
156
+ x = x + self.ffn(self.norm2(x))
157
+ return x, new_kv
158
+
159
+
160
+ # -------------------------------
161
+ # Главная модель GPTPyTorch (6 слоев)
162
+ # -------------------------------
163
+ class GPTPyTorch(nn.Module):
164
+ def __init__(self):
165
+ super().__init__()
166
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
167
+ self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
168
+ self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
169
+ self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
170
+ self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
171
+
172
+ signature = "Konstantin V Gbabko . original author © 2025"
173
+ bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
174
+ self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
175
+ self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
176
+
177
+ self.lm_head.weight = self.token_emb.weight
178
+ self.apply(self._init_weights)
179
+
180
+ def _init_weights(self, module):
181
+ if isinstance(module, nn.Linear):
182
+ # ИСПРАВЛЕНИЕ: Инициализация, масштабированная по глубине сети
183
+ std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
184
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
185
+ elif isinstance(module, nn.Embedding):
186
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
187
+ elif isinstance(module, nn.LayerNorm):
188
+ nn.init.zeros_(module.bias)
189
+ nn.init.ones_(module.weight)
190
+
191
+ def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
192
+ B, T = input_ids.shape
193
+ x = self.token_emb(input_ids)
194
+
195
+ pos_offset = 0
196
+ if past_kv is not None and past_kv[0] is not None:
197
+ pos_offset = past_kv[0][0].size(2)
198
+
199
+ x = self.pos_emb(x, pos_offset=pos_offset)
200
+
201
+ # Инициализация нового кеша
202
+ new_kv_cache = [] if past_kv is not None or T > 1 else None
203
+ current_past = past_kv
204
+
205
+ for i, block in enumerate(self.blocks):
206
+ layer_past = current_past[i] if (current_past and i < len(current_past)) else None
207
+ x, layer_kv = block(x, layer_past)
208
+
209
+ if new_kv_cache is not None:
210
+ new_kv_cache.append(layer_kv)
211
+
212
+ x = self.ln_f(x)
213
+ logits = self.lm_head(x)
214
+ return logits, new_kv_cache
215
+
216
+ @torch.no_grad()
217
+ def generate(
218
+ self,
219
+ input_ids: torch.Tensor,
220
+ max_new_tokens: int = 100,
221
+ temperature: float = 0.8,
222
+ top_p: float = 0.95,
223
+ repetition_penalty: float = 1.0,
224
+ do_sample: bool = True,
225
+ eos_token_id: int = 50256
226
+ ) -> torch.Tensor:
227
+ kv_cache = [None] * NUM_LAYERS
228
+ current_ids = input_ids.clone()
229
+
230
+ for step in range(max_new_tokens):
231
+ if step == 0:
232
+ input_for_model = current_ids
233
+ else:
234
+ input_for_model = current_ids[:, -1].unsqueeze(-1)
235
+
236
+ logits, kv_cache = self(input_for_model, kv_cache)
237
+ next_token_logits = logits[:, -1, :]
238
+
239
+ if temperature > 0:
240
+ next_token_logits = next_token_logits / temperature
241
+
242
+ # Repetition Penalty (логика сохранена)
243
+ if repetition_penalty != 1.0:
244
+ for i in range(current_ids.shape[0]):
245
+ unique_tokens = torch.unique(current_ids[i]).tolist()
246
+ for token_id in unique_tokens:
247
+ score = next_token_logits[i, token_id]
248
+ if score < 0:
249
+ next_token_logits[i, token_id] = score * repetition_penalty
250
+ else:
251
+ next_token_logits[i, token_id] = score / repetition_penalty
252
+
253
+ # Top-P сэмплирование (логика сохранена)
254
+ if do_sample and top_p < 1.0:
255
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
256
+ cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
257
+ sorted_indices_to_remove = cumulative_probs > top_p
258
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
259
+ sorted_indices_to_remove[:, 0] = False
260
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
261
+ next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))
262
+
263
+ # Сэмплирование
264
+ if do_sample and temperature > 0:
265
+ probs = torch.softmax(next_token_logits, dim=-1)
266
+ if torch.isnan(probs).any() or torch.isinf(probs).any():
267
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
268
+ else:
269
+ next_token = torch.multinomial(probs, num_samples=1)
270
+ else:
271
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
272
+
273
+ if next_token.item() == eos_token_id:
274
+ break
275
+
276
+ current_ids = torch.cat([current_ids, next_token], dim=1)
277
+
278
+ return current_ids
279
+
280
+
281
+ if __name__ == "__main__":
282
+ os.makedirs("models", exist_ok=True)
283
+
284
+ model = GPTPyTorch().to(device)
285
+ model.eval()
286
+
287
+ total_params = sum(p.numel() for p in model.parameters())
288
+ print(f"Device: {device}")
289
+ print(f"Total parameters: {total_params / 1e6:.2f}M") # ~124.44M
290
+
291
+ # 1. Проверка первого прохода (T=50)
292
+ input_ids_T50 = torch.randint(0, VOCAB_SIZE, (1, 50), device=device)
293
+ with torch.no_grad():
294
+ logits_50, kv_cache_50 = model(input_ids_T50)
295
+ expected_k_shape = (1, NUM_HEADS, 50, HEAD_DIM)
296
+ assert kv_cache_50[0][0].shape == expected_k_shape
297
+ print(f"Initial logits shape: {logits_50.shape}")
298
+ print(f"Initial KV-cache seqlen: {kv_cache_50[0][0].size(2)}")
299
+
300
+ # 2. Проверка инкрементального прохода (T=1)
301
+ input_ids_T1 = torch.randint(0, VOCAB_SIZE, (1, 1), device=device)
302
+ with torch.no_grad():
303
+ logits_51, kv_cache_51 = model(input_ids_T1, past_kv=kv_cache_50)
304
+
305
+ # Проверка длины кеша: 50 + 1 = 51
306
+ assert kv_cache_51[0][0].size(2) == 51
307
+ print(f"Incremental logits shape: {logits_51.shape}")
308
+ print(f"Incremental KV-cache seqlen: {kv_cache_51[0][0].size(2)}")
309
+
310
+ # 3. Проверка функции generate
311
+ generated = model.generate(input_ids_T50, max_new_tokens=5, temperature=0.8, top_p=0.9)
312
+ print(f"Generated sequence length (50 + 5): {generated.shape[1]}")
313
+
314
+ save_path = "models/JiRack_GPT_L6_PostNorm_fixed.pt"
315
+ torch.save(model.state_dict(), save_path)
316
+ print(f"Model successfully saved to {save_path}")
source/JiRack_H16_L24_V50257_D768_MSL8192_FF768x4.py ADDED
@@ -0,0 +1,318 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2025 CMS Manhattan
2
+ # All rights reserved.
3
+ # Author: Konstantin Vladimirovich Grabko
4
+ # Email: grabko@cmsmanhattan.com
5
+ # Phone: +1(516)777-0945
6
+ #
7
+ # This program is free software: you can redistribute it and/or modify
8
+ # it under the terms of the GNU General Public License as published by
9
+ # the Free Software Foundation, version 3 of the License.
10
+ #
11
+ # This program is distributed in the hope that it will be useful,
12
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
13
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
+ # GNU General Public License for more details.
15
+ #
16
+ # You should have received a copy of the GNU General Public License
17
+ # along with this program. If not, see <https://www.gnu.org/licenses/>.
18
+ #
19
+ # Additional terms:
20
+ # Any commercial use or distribution of this software or derivative works
21
+ # requires explicit written permission from the copyright holder.
22
+
23
+ import os
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from typing import Optional, Tuple, List
28
+ import math
29
+
30
+ # ========================================
31
+ # Model Configuration (L=24, H=16)
32
+ # ========================================
33
+ VOCAB_SIZE = 50257
34
+ MODEL_DIM = 768
35
+ NUM_HEADS = 16 # Changed to 16
36
+ NUM_LAYERS = 24 # Changed to 24
37
+ MAX_SEQ_LEN = 8192
38
+ FFN_HIDDEN_DIM = 4 * MODEL_DIM
39
+ HEAD_DIM = MODEL_DIM // NUM_HEADS # Recalculated: 768 / 16 = 48
40
+
41
+ # ИСПРАВЛЕНИЕ: ROCm/HIP-совместимая проверка устройства
42
+ if torch.cuda.is_available():
43
+ device = torch.device("cuda")
44
+ elif hasattr(torch, 'hip') and torch.hip.is_available():
45
+ device = torch.device("cuda")
46
+ else:
47
+ device = torch.device("cpu")
48
+
49
+ # -------------------------------
50
+ # Learned Positional Embedding
51
+ # -------------------------------
52
+ class LearnedPositionalEmbedding(nn.Module):
53
+ def __init__(self, max_seq_len: int, embed_dim: int):
54
+ super().__init__()
55
+ self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
56
+
57
+ def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
58
+ seq_len = x.size(1)
59
+ if pos_offset + seq_len > self.pos_emb.size(0):
60
+ raise ValueError("Sequence length exceeds MAX_SEQ_LEN defined in position embedding.")
61
+
62
+ pos = self.pos_emb[pos_offset : pos_offset + seq_len]
63
+ return x + pos.unsqueeze(0)
64
+
65
+
66
+ # -------------------------------
67
+ # MultiHeadAttention (MHA)
68
+ # -------------------------------
69
+ class MultiHeadAttention(nn.Module):
70
+ def __init__(self):
71
+ super().__init__()
72
+ self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
73
+ self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
74
+ self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
75
+ self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
76
+ self.scale = HEAD_DIM ** -0.5
77
+
78
+ def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
79
+ B, T, D = x.shape
80
+
81
+ # NOTE: D is MODEL_DIM
82
+ q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
83
+ k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
84
+ v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
85
+
86
+ # 1. KV-кеш и определение смещения
87
+ pos_offset = 0
88
+ seqlen_k_new = k.size(2)
89
+ if past_kv is not None:
90
+ past_k, past_v = past_kv
91
+ k = torch.cat([past_k, k], dim=2)
92
+ v = torch.cat([past_v, v], dim=2)
93
+ pos_offset = past_k.size(2)
94
+
95
+ seqlen_k = k.size(2) # Общая длина K
96
+ new_kv = (k, v)
97
+
98
+ # 2. Расчет внимания
99
+ attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
100
+
101
+ # 3. КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ МАСКИРОВАНИЯ (Causal Mask)
102
+ if T == seqlen_k_new and seqlen_k > 0:
103
+ # Создаем маску T (query length) x seqlen_k (key length)
104
+ mask = torch.full((T, seqlen_k),
105
+ float("-inf"),
106
+ device=x.device,
107
+ dtype=attn.dtype)
108
+
109
+ # Разрешаем всем новым токенам видеть все старые токены (past_kv)
110
+ mask[:, :pos_offset] = 0.0
111
+
112
+ # Применяем треугольную маску для текущих T токенов
113
+ current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
114
+ mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
115
+
116
+ # Применяем маску к весам внимания
117
+ attn = attn + mask[None, None, :, :]
118
+
119
+ # 4. Выход
120
+ attn = F.softmax(attn, dim=-1)
121
+ out = torch.matmul(attn, v)
122
+ out = out.transpose(1, 2).contiguous().view(B, T, D)
123
+ out = self.out_proj(out)
124
+
125
+ return out, new_kv
126
+
127
+
128
+ # -------------------------------
129
+ # FeedForward (GELU, GPT-style)
130
+ # -------------------------------
131
+ class FeedForward(nn.Module):
132
+ def __init__(self):
133
+ super().__init__()
134
+ self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
135
+ self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
136
+
137
+ def forward(self, x):
138
+ return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
139
+
140
+
141
+ # -------------------------------
142
+ # Transformer Block (Post-Norm, GPT-style)
143
+ # -------------------------------
144
+ class TransformerBlock(nn.Module):
145
+ def __init__(self):
146
+ super().__init__()
147
+ self.attn = MultiHeadAttention()
148
+ self.ffn = FeedForward()
149
+ # ИСПРАВЛЕНИЕ: Добавлен стандартный eps
150
+ self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
151
+ self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
152
+
153
+ def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
154
+ # Post-Normalization (GPT Style)
155
+ attn_out, new_kv = self.attn(self.norm1(x), past_kv)
156
+ x = x + attn_out
157
+ x = x + self.ffn(self.norm2(x))
158
+ return x, new_kv
159
+
160
+
161
+ # -------------------------------
162
+ # Главная модель GPTPyTorch (L=24, H=16)
163
+ # -------------------------------
164
+ class GPTPyTorch(nn.Module):
165
+ def __init__(self):
166
+ super().__init__()
167
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
168
+ self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
169
+ self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
170
+ self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
171
+ self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
172
+
173
+ signature = "Konstantin V Gbabko . original author © 2025"
174
+ bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
175
+ self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
176
+ self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
177
+
178
+ self.lm_head.weight = self.token_emb.weight
179
+ self.apply(self._init_weights)
180
+
181
+ def _init_weights(self, module):
182
+ if isinstance(module, nn.Linear):
183
+ # ИСПРАВЛЕНИЕ: Инициализация, масштабированная по глубине сети
184
+ std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
185
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
186
+ elif isinstance(module, nn.Embedding):
187
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
188
+ elif isinstance(module, nn.LayerNorm):
189
+ nn.init.zeros_(module.bias)
190
+ nn.init.ones_(module.weight)
191
+
192
+ def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
193
+ B, T = input_ids.shape
194
+ x = self.token_emb(input_ids)
195
+
196
+ pos_offset = 0
197
+ if past_kv is not None and past_kv[0] is not None:
198
+ pos_offset = past_kv[0][0].size(2)
199
+
200
+ x = self.pos_emb(x, pos_offset=pos_offset)
201
+
202
+ # Инициализация нового кеша
203
+ new_kv_cache = [] if past_kv is not None or T > 1 else None
204
+ current_past = past_kv
205
+
206
+ for i, block in enumerate(self.blocks):
207
+ layer_past = current_past[i] if (current_past and i < len(current_past)) else None
208
+ x, layer_kv = block(x, layer_past)
209
+
210
+ if new_kv_cache is not None:
211
+ new_kv_cache.append(layer_kv)
212
+
213
+ x = self.ln_f(x)
214
+ logits = self.lm_head(x)
215
+ return logits, new_kv_cache
216
+
217
+ @torch.no_grad()
218
+ def generate(
219
+ self,
220
+ input_ids: torch.Tensor,
221
+ max_new_tokens: int = 100,
222
+ temperature: float = 0.8,
223
+ top_p: float = 0.95,
224
+ repetition_penalty: float = 1.0,
225
+ do_sample: bool = True,
226
+ eos_token_id: int = 50256
227
+ ) -> torch.Tensor:
228
+ kv_cache = [None] * NUM_LAYERS
229
+ current_ids = input_ids.clone()
230
+
231
+ for step in range(max_new_tokens):
232
+ if step == 0:
233
+ input_for_model = current_ids
234
+ else:
235
+ input_for_model = current_ids[:, -1].unsqueeze(-1)
236
+
237
+ logits, kv_cache = self(input_for_model, kv_cache)
238
+ next_token_logits = logits[:, -1, :]
239
+
240
+ if temperature > 0:
241
+ next_token_logits = next_token_logits / temperature
242
+
243
+ # Repetition Penalty (логика сохранена)
244
+ if repetition_penalty != 1.0:
245
+ for i in range(current_ids.shape[0]):
246
+ unique_tokens = torch.unique(current_ids[i]).tolist()
247
+ for token_id in unique_tokens:
248
+ score = next_token_logits[i, token_id]
249
+ if score < 0:
250
+ next_token_logits[i, token_id] = score * repetition_penalty
251
+ else:
252
+ next_token_logits[i, token_id] = score / repetition_penalty
253
+
254
+ # Top-P сэмплирование (логика сохранена)
255
+ if do_sample and top_p < 1.0:
256
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
257
+ cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
258
+ sorted_indices_to_remove = cumulative_probs > top_p
259
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
260
+ sorted_indices_to_remove[:, 0] = False
261
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
262
+ next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))
263
+
264
+ # Сэмплирование
265
+ if do_sample and temperature > 0:
266
+ probs = torch.softmax(next_token_logits, dim=-1)
267
+ if torch.isnan(probs).any() or torch.isinf(probs).any():
268
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
269
+ else:
270
+ next_token = torch.multinomial(probs, num_samples=1)
271
+ else:
272
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
273
+
274
+ if next_token.item() == eos_token_id:
275
+ break
276
+
277
+ current_ids = torch.cat([current_ids, next_token], dim=1)
278
+
279
+ return current_ids
280
+
281
+
282
+ if __name__ == "__main__":
283
+ os.makedirs("models", exist_ok=True)
284
+
285
+ model = GPTPyTorch().to(device)
286
+ model.eval()
287
+
288
+ total_params = sum(p.numel() for p in model.parameters())
289
+ print(f"Device: {device}")
290
+ print(f"Total parameters: {total_params / 1e6:.2f}M") # ~221.75M
291
+
292
+ # 1. Проверка первого прохода (T=50)
293
+ input_ids_T50 = torch.randint(0, VOCAB_SIZE, (1, 50), device=device)
294
+ with torch.no_grad():
295
+ logits_50, kv_cache_50 = model(input_ids_T50)
296
+ expected_k_shape = (1, NUM_HEADS, 50, HEAD_DIM)
297
+ # Проверяем, что размерность головы HEAD_DIM (48) соблюдена
298
+ assert kv_cache_50[0][0].shape == expected_k_shape
299
+ print(f"Initial logits shape: {logits_50.shape}")
300
+ print(f"Initial KV-cache seqlen: {kv_cache_50[0][0].size(2)} (Head Dim: {kv_cache_50[0][0].size(3)})")
301
+
302
+ # 2. Проверка инкрементального прохода (T=1)
303
+ input_ids_T1 = torch.randint(0, VOCAB_SIZE, (1, 1), device=device)
304
+ with torch.no_grad():
305
+ logits_51, kv_cache_51 = model(input_ids_T1, past_kv=kv_cache_50)
306
+
307
+ # Проверка длины кеша: 50 + 1 = 51
308
+ assert kv_cache_51[0][0].size(2) == 51
309
+ print(f"Incremental logits shape: {logits_51.shape}")
310
+ print(f"Incremental KV-cache seqlen: {kv_cache_51[0][0].size(2)}")
311
+
312
+ # 3. Проверка функции generate
313
+ generated = model.generate(input_ids_T50, max_new_tokens=5, temperature=0.8, top_p=0.9)
314
+ print(f"Generated sequence length (50 + 5): {generated.shape[1]}")
315
+
316
+ save_path = "models/JiRack_GPT_L24_H16_PostNorm_fixed.pt"
317
+ torch.save(model.state_dict(), save_path)
318
+ print(f"Model successfully saved to {save_path}")
source/JiRack_H16_L32_V50257_D768_MSL8192_FF768x4.py ADDED
@@ -0,0 +1,318 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2025 CMS Manhattan
2
+ # All rights reserved.
3
+ # Author: Konstantin Vladimirovich Grabko
4
+ # Email: grabko@cmsmanhattan.com
5
+ # Phone: +1(516)777-0945
6
+ #
7
+ # This program is free software: you can redistribute it and/or modify
8
+ # it under the terms of the GNU General Public License as published by
9
+ # the Free Software Foundation, version 3 of the License.
10
+ #
11
+ # This program is distributed in the hope that it will be useful,
12
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
13
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
+ # GNU General Public License for more details.
15
+ #
16
+ # You should have received a copy of the GNU General Public License
17
+ # along with this program. If not, see <https://www.gnu.org/licenses/>.
18
+ #
19
+ # Additional terms:
20
+ # Any commercial use or distribution of this software or derivative works
21
+ # requires explicit written permission from the copyright holder.
22
+
23
+ import os
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from typing import Optional, Tuple, List
28
+ import math
29
+
30
+ # ========================================
31
+ # Model Configuration (L=32, H=16, D=768)
32
+ # ========================================
33
+ VOCAB_SIZE = 50257
34
+ MODEL_DIM = 768
35
+ NUM_HEADS = 16
36
+ NUM_LAYERS = 32 # Deepest version yet
37
+ MAX_SEQ_LEN = 8192
38
+ FFN_HIDDEN_DIM = 4 * MODEL_DIM
39
+ HEAD_DIM = MODEL_DIM // NUM_HEADS # 768 / 16 = 48
40
+
41
+ # ИСПРАВЛЕНИЕ: ROCm/HIP-совместимая проверка устройства
42
+ if torch.cuda.is_available():
43
+ device = torch.device("cuda")
44
+ elif hasattr(torch, 'hip') and torch.hip.is_available():
45
+ device = torch.device("cuda")
46
+ else:
47
+ device = torch.device("cpu")
48
+
49
+ # -------------------------------
50
+ # Learned Positional Embedding
51
+ # -------------------------------
52
+ class LearnedPositionalEmbedding(nn.Module):
53
+ def __init__(self, max_seq_len: int, embed_dim: int):
54
+ super().__init__()
55
+ self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
56
+
57
+ def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
58
+ seq_len = x.size(1)
59
+ if pos_offset + seq_len > self.pos_emb.size(0):
60
+ raise ValueError("Sequence length exceeds MAX_SEQ_LEN defined in position embedding.")
61
+
62
+ pos = self.pos_emb[pos_offset : pos_offset + seq_len]
63
+ return x + pos.unsqueeze(0)
64
+
65
+
66
+ # -------------------------------
67
+ # MultiHeadAttention (MHA)
68
+ # -------------------------------
69
+ class MultiHeadAttention(nn.Module):
70
+ def __init__(self):
71
+ super().__init__()
72
+ self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
73
+ self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
74
+ self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
75
+ self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
76
+ self.scale = HEAD_DIM ** -0.5
77
+
78
+ def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
79
+ B, T, D = x.shape
80
+
81
+ q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
82
+ k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
83
+ v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
84
+
85
+ # 1. KV-кеш и определение смещения
86
+ pos_offset = 0
87
+ seqlen_k_new = k.size(2)
88
+ if past_kv is not None:
89
+ past_k, past_v = past_kv
90
+ k = torch.cat([past_k, k], dim=2)
91
+ v = torch.cat([past_v, v], dim=2)
92
+ pos_offset = past_k.size(2)
93
+
94
+ seqlen_k = k.size(2) # Общая длина K
95
+ new_kv = (k, v)
96
+
97
+ # 2. Расчет внимания
98
+ attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
99
+
100
+ # 3. КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ МАСКИРОВАНИЯ (Causal Mask)
101
+ if T == seqlen_k_new and seqlen_k > 0:
102
+ # Создаем маску T (query length) x seqlen_k (key length)
103
+ mask = torch.full((T, seqlen_k),
104
+ float("-inf"),
105
+ device=x.device,
106
+ dtype=attn.dtype)
107
+
108
+ # Разрешаем всем новым токенам видеть все старые токены (past_kv)
109
+ mask[:, :pos_offset] = 0.0
110
+
111
+ # Применяем треугольную маску для текущих T токенов
112
+ current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
113
+ mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
114
+
115
+ # Применяем маску к весам внимания
116
+ attn = attn + mask[None, None, :, :]
117
+
118
+ # 4. Выход
119
+ attn = F.softmax(attn, dim=-1)
120
+ out = torch.matmul(attn, v)
121
+ out = out.transpose(1, 2).contiguous().view(B, T, D)
122
+ out = self.out_proj(out)
123
+
124
+ return out, new_kv
125
+
126
+
127
+ # -------------------------------
128
+ # FeedForward (GELU, GPT-style)
129
+ # -------------------------------
130
+ class FeedForward(nn.Module):
131
+ def __init__(self):
132
+ super().__init__()
133
+ self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
134
+ self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
135
+
136
+ def forward(self, x):
137
+ return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
138
+
139
+
140
+ # -------------------------------
141
+ # Transformer Block (Post-Norm, GPT-style)
142
+ # -------------------------------
143
+ class TransformerBlock(nn.Module):
144
+ def __init__(self):
145
+ super().__init__()
146
+ self.attn = MultiHeadAttention()
147
+ self.ffn = FeedForward()
148
+ # ИСПРАВЛЕНИЕ: Добавлен стандартный eps
149
+ self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
150
+ self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
151
+
152
+ def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
153
+ # Post-Normalization (GPT Style)
154
+ attn_out, new_kv = self.attn(self.norm1(x), past_kv)
155
+ x = x + attn_out
156
+ x = x + self.ffn(self.norm2(x))
157
+ return x, new_kv
158
+
159
+
160
+ # -------------------------------
161
+ # Главная модель GPTPyTorch (L=32, H=16)
162
+ # -------------------------------
163
+ class GPTPyTorch(nn.Module):
164
+ def __init__(self):
165
+ super().__init__()
166
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
167
+ self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
168
+ self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
169
+ self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
170
+ self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
171
+
172
+ signature = "Konstantin V Gbabko . original author © 2025"
173
+ bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
174
+ self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
175
+ self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
176
+
177
+ self.lm_head.weight = self.token_emb.weight
178
+ self.apply(self._init_weights)
179
+
180
+ def _init_weights(self, module):
181
+ if isinstance(module, nn.Linear):
182
+ # ИСПРАВЛЕНИЕ: Инициализация, масштабированная по глубине сети.
183
+ # Critical for NUM_LAYERS = 32
184
+ std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
185
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
186
+ elif isinstance(module, nn.Embedding):
187
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
188
+ elif isinstance(module, nn.LayerNorm):
189
+ nn.init.zeros_(module.bias)
190
+ nn.init.ones_(module.weight)
191
+
192
+ def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
193
+ B, T = input_ids.shape
194
+ x = self.token_emb(input_ids)
195
+
196
+ pos_offset = 0
197
+ if past_kv is not None and past_kv[0] is not None:
198
+ pos_offset = past_kv[0][0].size(2)
199
+
200
+ x = self.pos_emb(x, pos_offset=pos_offset)
201
+
202
+ # Инициализация нового кеша
203
+ new_kv_cache = [] if past_kv is not None or T > 1 else None
204
+ current_past = past_kv
205
+
206
+ for i, block in enumerate(self.blocks):
207
+ layer_past = current_past[i] if (current_past and i < len(current_past)) else None
208
+ x, layer_kv = block(x, layer_past)
209
+
210
+ if new_kv_cache is not None:
211
+ new_kv_cache.append(layer_kv)
212
+
213
+ x = self.ln_f(x)
214
+ logits = self.lm_head(x)
215
+ return logits, new_kv_cache
216
+
217
+ @torch.no_grad()
218
+ def generate(
219
+ self,
220
+ input_ids: torch.Tensor,
221
+ max_new_tokens: int = 100,
222
+ temperature: float = 0.8,
223
+ top_p: float = 0.95,
224
+ repetition_penalty: float = 1.0,
225
+ do_sample: bool = True,
226
+ eos_token_id: int = 50256
227
+ ) -> torch.Tensor:
228
+ kv_cache = [None] * NUM_LAYERS
229
+ current_ids = input_ids.clone()
230
+
231
+ for step in range(max_new_tokens):
232
+ if step == 0:
233
+ input_for_model = current_ids
234
+ else:
235
+ input_for_model = current_ids[:, -1].unsqueeze(-1)
236
+
237
+ logits, kv_cache = self(input_for_model, kv_cache)
238
+ next_token_logits = logits[:, -1, :]
239
+
240
+ if temperature > 0:
241
+ next_token_logits = next_token_logits / temperature
242
+
243
+ # Repetition Penalty (логика сохранена)
244
+ if repetition_penalty != 1.0:
245
+ for i in range(current_ids.shape[0]):
246
+ unique_tokens = torch.unique(current_ids[i]).tolist()
247
+ for token_id in unique_tokens:
248
+ score = next_token_logits[i, token_id]
249
+ if score < 0:
250
+ next_token_logits[i, token_id] = score * repetition_penalty
251
+ else:
252
+ next_token_logits[i, token_id] = score / repetition_penalty
253
+
254
+ # Top-P сэмплирование (логика сохранена)
255
+ if do_sample and top_p < 1.0:
256
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
257
+ cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
258
+ sorted_indices_to_remove = cumulative_probs > top_p
259
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
260
+ sorted_indices_to_remove[:, 0] = False
261
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
262
+ next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))
263
+
264
+ # Сэмплирование
265
+ if do_sample and temperature > 0:
266
+ probs = torch.softmax(next_token_logits, dim=-1)
267
+ if torch.isnan(probs).any() or torch.isinf(probs).any():
268
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
269
+ else:
270
+ next_token = torch.multinomial(probs, num_samples=1)
271
+ else:
272
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
273
+
274
+ if next_token.item() == eos_token_id:
275
+ break
276
+
277
+ current_ids = torch.cat([current_ids, next_token], dim=1)
278
+
279
+ return current_ids
280
+
281
+
282
+ if __name__ == "__main__":
283
+ os.makedirs("models", exist_ok=True)
284
+
285
+ model = GPTPyTorch().to(device)
286
+ model.eval()
287
+
288
+ total_params = sum(p.numel() for p in model.parameters())
289
+ print(f"Device: {device}")
290
+ print(f"Total parameters: {total_params / 1e6:.2f}M") # ~295.1M
291
+
292
+ # 1. Проверка первого прохода (T=50)
293
+ input_ids_T50 = torch.randint(0, VOCAB_SIZE, (1, 50), device=device)
294
+ with torch.no_grad():
295
+ logits_50, kv_cache_50 = model(input_ids_T50)
296
+ expected_k_shape = (1, NUM_HEADS, 50, HEAD_DIM)
297
+ # Проверяем, что размерность головы HEAD_DIM (48) соблюдена
298
+ assert kv_cache_50[0][0].shape == expected_k_shape
299
+ print(f"Initial logits shape: {logits_50.shape}")
300
+ print(f"Initial KV-cache seqlen: {kv_cache_50[0][0].size(2)} (Head Dim: {kv_cache_50[0][0].size(3)})")
301
+
302
+ # 2. Проверка инкрементального прохода (T=1)
303
+ input_ids_T1 = torch.randint(0, VOCAB_SIZE, (1, 1), device=device)
304
+ with torch.no_grad():
305
+ logits_51, kv_cache_51 = model(input_ids_T1, past_kv=kv_cache_50)
306
+
307
+ # Проверка длины кеша: 50 + 1 = 51
308
+ assert kv_cache_51[0][0].size(2) == 51
309
+ print(f"Incremental logits shape: {logits_51.shape}")
310
+ print(f"Incremental KV-cache seqlen: {kv_cache_51[0][0].size(2)}")
311
+
312
+ # 3. Проверка функции generate
313
+ generated = model.generate(input_ids_T50, max_new_tokens=5, temperature=0.8, top_p=0.9)
314
+ print(f"Generated sequence length (50 + 5): {generated.shape[1]}")
315
+
316
+ save_path = "models/JiRack_GPT_L32_H16_PostNorm_fixed.pt"
317
+ torch.save(model.state_dict(), save_path)
318
+ print(f"Model successfully saved to {save_path}")
source/JiRack_H6_L6_V50257_D768_MSL8192_FF768x4.py ADDED
@@ -0,0 +1,317 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2025 CMS Manhattan
2
+ # All rights reserved.
3
+ # Author: Konstantin Vladimirovich Grabko
4
+ # Email: grabko@cmsmanhattan.com
5
+ # Phone: +1(516)777-0945
6
+ #
7
+ # This program is free software: you can redistribute it and/or modify
8
+ # it under the terms of the GNU General Public License as published by
9
+ # the Free Software Foundation, version 3 of the License.
10
+ #
11
+ # This program is distributed in the hope that it will be useful,
12
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
13
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
+ # GNU General Public License for more details.
15
+ #
16
+ # You should have received a copy of the GNU General Public License
17
+ # along with this program. If not, see <https://www.gnu.org/licenses/>.
18
+ #
19
+ # Additional terms:
20
+ # Any commercial use or distribution of this software or derivative works
21
+ # requires explicit written permission from the copyright holder.
22
+
23
+ import os
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from typing import Optional, Tuple, List
28
+ import math
29
+
30
+ # ========================================
31
+ # Model Configuration (L=6, H=6, D=768)
32
+ # ========================================
33
+ VOCAB_SIZE = 50257
34
+ MODEL_DIM = 768
35
+ NUM_HEADS = 6 # Changed to 6
36
+ NUM_LAYERS = 6 # Set to 6 layers
37
+ MAX_SEQ_LEN = 8192
38
+ FFN_HIDDEN_DIM = 4 * MODEL_DIM
39
+ HEAD_DIM = MODEL_DIM // NUM_HEADS # 768 / 6 = 128
40
+
41
+ # ИСПРАВЛЕНИЕ: ROCm/HIP-совместимая проверка устройства
42
+ if torch.cuda.is_available():
43
+ device = torch.device("cuda")
44
+ elif hasattr(torch, 'hip') and torch.hip.is_available():
45
+ device = torch.device("cuda")
46
+ else:
47
+ device = torch.device("cpu")
48
+
49
+ # -------------------------------
50
+ # Learned Positional Embedding
51
+ # -------------------------------
52
+ class LearnedPositionalEmbedding(nn.Module):
53
+ def __init__(self, max_seq_len: int, embed_dim: int):
54
+ super().__init__()
55
+ self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
56
+
57
+ def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
58
+ seq_len = x.size(1)
59
+ if pos_offset + seq_len > self.pos_emb.size(0):
60
+ raise ValueError("Sequence length exceeds MAX_SEQ_LEN defined in position embedding.")
61
+
62
+ pos = self.pos_emb[pos_offset : pos_offset + seq_len]
63
+ return x + pos.unsqueeze(0)
64
+
65
+
66
+ # -------------------------------
67
+ # MultiHeadAttention (MHA)
68
+ # -------------------------------
69
+ class MultiHeadAttention(nn.Module):
70
+ def __init__(self):
71
+ super().__init__()
72
+ self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
73
+ self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
74
+ self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
75
+ self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
76
+ self.scale = HEAD_DIM ** -0.5
77
+
78
+ def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
79
+ B, T, D = x.shape
80
+
81
+ q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
82
+ k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
83
+ v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
84
+
85
+ # 1. KV-кеш и определение смещения
86
+ pos_offset = 0
87
+ seqlen_k_new = k.size(2)
88
+ if past_kv is not None:
89
+ past_k, past_v = past_kv
90
+ k = torch.cat([past_k, k], dim=2)
91
+ v = torch.cat([past_v, v], dim=2)
92
+ pos_offset = past_k.size(2)
93
+
94
+ seqlen_k = k.size(2) # Общая длина K
95
+ new_kv = (k, v)
96
+
97
+ # 2. Расчет внимания
98
+ attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
99
+
100
+ # 3. КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ МАСКИРОВАНИЯ (Causal Mask)
101
+ if T == seqlen_k_new and seqlen_k > 0:
102
+ # Создаем маску T (query length) x seqlen_k (key length)
103
+ mask = torch.full((T, seqlen_k),
104
+ float("-inf"),
105
+ device=x.device,
106
+ dtype=attn.dtype)
107
+
108
+ # Разрешаем всем новым токенам видеть все старые токены (past_kv)
109
+ mask[:, :pos_offset] = 0.0
110
+
111
+ # Применяем треугольную маску для текущих T токенов
112
+ current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
113
+ mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
114
+
115
+ # Применяем маску к весам внимания
116
+ attn = attn + mask[None, None, :, :]
117
+
118
+ # 4. Выход
119
+ attn = F.softmax(attn, dim=-1)
120
+ out = torch.matmul(attn, v)
121
+ out = out.transpose(1, 2).contiguous().view(B, T, D)
122
+ out = self.out_proj(out)
123
+
124
+ return out, new_kv
125
+
126
+
127
+ # -------------------------------
128
+ # FeedForward (GELU, GPT-style)
129
+ # -------------------------------
130
+ class FeedForward(nn.Module):
131
+ def __init__(self):
132
+ super().__init__()
133
+ self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
134
+ self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
135
+
136
+ def forward(self, x):
137
+ return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
138
+
139
+
140
+ # -------------------------------
141
+ # Transformer Block (Post-Norm, GPT-style)
142
+ # -------------------------------
143
+ class TransformerBlock(nn.Module):
144
+ def __init__(self):
145
+ super().__init__()
146
+ self.attn = MultiHeadAttention()
147
+ self.ffn = FeedForward()
148
+ # ИСПРАВЛЕНИЕ: Добавлен стандартный eps
149
+ self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
150
+ self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
151
+
152
+ def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
153
+ # Post-Normalization (GPT Style)
154
+ attn_out, new_kv = self.attn(self.norm1(x), past_kv)
155
+ x = x + attn_out
156
+ x = x + self.ffn(self.norm2(x))
157
+ return x, new_kv
158
+
159
+
160
+ # -------------------------------
161
+ # Главная модель GPTPyTorch (L=6, H=6)
162
+ # -------------------------------
163
+ class GPTPyTorch(nn.Module):
164
+ def __init__(self):
165
+ super().__init__()
166
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
167
+ self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
168
+ self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
169
+ self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
170
+ self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
171
+
172
+ signature = "Konstantin V Gbabko . original author © 2025"
173
+ bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
174
+ self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
175
+ self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
176
+
177
+ self.lm_head.weight = self.token_emb.weight
178
+ self.apply(self._init_weights)
179
+
180
+ def _init_weights(self, module):
181
+ if isinstance(module, nn.Linear):
182
+ # ИСПРАВЛЕНИЕ: Инициализация, масштабированная по глубине сети.
183
+ std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
184
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
185
+ elif isinstance(module, nn.Embedding):
186
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
187
+ elif isinstance(module, nn.LayerNorm):
188
+ nn.init.zeros_(module.bias)
189
+ nn.init.ones_(module.weight)
190
+
191
+ def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
192
+ B, T = input_ids.shape
193
+ x = self.token_emb(input_ids)
194
+
195
+ pos_offset = 0
196
+ if past_kv is not None and past_kv[0] is not None:
197
+ pos_offset = past_kv[0][0].size(2)
198
+
199
+ x = self.pos_emb(x, pos_offset=pos_offset)
200
+
201
+ # Инициализация нового кеша
202
+ new_kv_cache = [] if past_kv is not None or T > 1 else None
203
+ current_past = past_kv
204
+
205
+ for i, block in enumerate(self.blocks):
206
+ layer_past = current_past[i] if (current_past and i < len(current_past)) else None
207
+ x, layer_kv = block(x, layer_past)
208
+
209
+ if new_kv_cache is not None:
210
+ new_kv_cache.append(layer_kv)
211
+
212
+ x = self.ln_f(x)
213
+ logits = self.lm_head(x)
214
+ return logits, new_kv_cache
215
+
216
+ @torch.no_grad()
217
+ def generate(
218
+ self,
219
+ input_ids: torch.Tensor,
220
+ max_new_tokens: int = 100,
221
+ temperature: float = 0.8,
222
+ top_p: float = 0.95,
223
+ repetition_penalty: float = 1.0,
224
+ do_sample: bool = True,
225
+ eos_token_id: int = 50256
226
+ ) -> torch.Tensor:
227
+ kv_cache = [None] * NUM_LAYERS
228
+ current_ids = input_ids.clone()
229
+
230
+ for step in range(max_new_tokens):
231
+ if step == 0:
232
+ input_for_model = current_ids
233
+ else:
234
+ input_for_model = current_ids[:, -1].unsqueeze(-1)
235
+
236
+ logits, kv_cache = self(input_for_model, kv_cache)
237
+ next_token_logits = logits[:, -1, :]
238
+
239
+ if temperature > 0:
240
+ next_token_logits = next_token_logits / temperature
241
+
242
+ # Repetition Penalty (логика сохранена)
243
+ if repetition_penalty != 1.0:
244
+ for i in range(current_ids.shape[0]):
245
+ unique_tokens = torch.unique(current_ids[i]).tolist()
246
+ for token_id in unique_tokens:
247
+ score = next_token_logits[i, token_id]
248
+ if score < 0:
249
+ next_token_logits[i, token_id] = score * repetition_penalty
250
+ else:
251
+ next_token_logits[i, token_id] = score / repetition_penalty
252
+
253
+ # Top-P сэмплирование (логика сохранена)
254
+ if do_sample and top_p < 1.0:
255
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
256
+ cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
257
+ sorted_indices_to_remove = cumulative_probs > top_p
258
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
259
+ sorted_indices_to_remove[:, 0] = False
260
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
261
+ next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))
262
+
263
+ # Сэмплирование
264
+ if do_sample and temperature > 0:
265
+ probs = torch.softmax(next_token_logits, dim=-1)
266
+ if torch.isnan(probs).any() or torch.isinf(probs).any():
267
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
268
+ else:
269
+ next_token = torch.multinomial(probs, num_samples=1)
270
+ else:
271
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
272
+
273
+ if next_token.item() == eos_token_id:
274
+ break
275
+
276
+ current_ids = torch.cat([current_ids, next_token], dim=1)
277
+
278
+ return current_ids
279
+
280
+
281
+ if __name__ == "__main__":
282
+ os.makedirs("models", exist_ok=True)
283
+
284
+ model = GPTPyTorch().to(device)
285
+ model.eval()
286
+
287
+ total_params = sum(p.numel() for p in model.parameters())
288
+ print(f"Device: {device}")
289
+ print(f"Total parameters: {total_params / 1e6:.2f}M") # ~124.44M
290
+
291
+ # 1. Проверка первого прохода (T=50)
292
+ input_ids_T50 = torch.randint(0, VOCAB_SIZE, (1, 50), device=device)
293
+ with torch.no_grad():
294
+ logits_50, kv_cache_50 = model(input_ids_T50)
295
+ expected_k_shape = (1, NUM_HEADS, 50, HEAD_DIM)
296
+ # Проверяем, что размерность головы HEAD_DIM (128) соблюдена
297
+ assert kv_cache_50[0][0].shape == expected_k_shape
298
+ print(f"Initial logits shape: {logits_50.shape}")
299
+ print(f"Initial KV-cache seqlen: {kv_cache_50[0][0].size(2)} (Head Dim: {kv_cache_50[0][0].size(3)})")
300
+
301
+ # 2. Проверка инкрементального прохода (T=1)
302
+ input_ids_T1 = torch.randint(0, VOCAB_SIZE, (1, 1), device=device)
303
+ with torch.no_grad():
304
+ logits_51, kv_cache_51 = model(input_ids_T1, past_kv=kv_cache_50)
305
+
306
+ # Проверка длины кеша: 50 + 1 = 51
307
+ assert kv_cache_51[0][0].size(2) == 51
308
+ print(f"Incremental logits shape: {logits_51.shape}")
309
+ print(f"Incremental KV-cache seqlen: {kv_cache_51[0][0].size(2)}")
310
+
311
+ # 3. Проверка функции generate
312
+ generated = model.generate(input_ids_T50, max_new_tokens=5, temperature=0.8, top_p=0.9)
313
+ print(f"Generated sequence length (50 + 5): {generated.shape[1]}")
314
+
315
+ save_path = "models/JiRack_GPT_L6_H6_PostNorm_fixed.pt"
316
+ torch.save(model.state_dict(), save_path)
317
+ print(f"Model successfully saved to {save_path}")
source/JiRack_H8_L6_V50257_D768_MSL8192_FF768x4.py ADDED
@@ -0,0 +1,317 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2025 CMS Manhattan
2
+ # All rights reserved.
3
+ # Author: Konstantin Vladimirovich Grabko
4
+ # Email: grabko@cmsmanhattan.com
5
+ # Phone: +1(516)777-0945
6
+ #
7
+ # This program is free software: you can redistribute it and/or modify
8
+ # it under the terms of the GNU General Public License as published by
9
+ # the Free Software Foundation, version 3 of the License.
10
+ #
11
+ # This program is distributed in the hope that it will be useful,
12
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
13
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
+ # GNU General Public License for more details.
15
+ #
16
+ # You should have received a copy of the GNU General Public License
17
+ # along with this program. If not, see <https://www.gnu.org/licenses/>.
18
+ #
19
+ # Additional terms:
20
+ # Any commercial use or distribution of this software or derivative works
21
+ # requires explicit written permission from the copyright holder.
22
+
23
+ import os
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from typing import Optional, Tuple, List
28
+ import math
29
+
30
+ # ========================================
31
+ # Model Configuration (L=6, H=8, D=768)
32
+ # ========================================
33
+ VOCAB_SIZE = 50257
34
+ MODEL_DIM = 768
35
+ NUM_HEADS = 8 # Set to 8
36
+ NUM_LAYERS = 6
37
+ MAX_SEQ_LEN = 8192
38
+ FFN_HIDDEN_DIM = 4 * MODEL_DIM
39
+ HEAD_DIM = MODEL_DIM // NUM_HEADS # Recalculated: 768 / 8 = 96
40
+
41
+ # ИСПРАВЛЕНИЕ: ROCm/HIP-совместимая проверка устройства
42
+ if torch.cuda.is_available():
43
+ device = torch.device("cuda")
44
+ elif hasattr(torch, 'hip') and torch.hip.is_available():
45
+ device = torch.device("cuda")
46
+ else:
47
+ device = torch.device("cpu")
48
+
49
+ # -------------------------------
50
+ # Learned Positional Embedding
51
+ # -------------------------------
52
+ class LearnedPositionalEmbedding(nn.Module):
53
+ def __init__(self, max_seq_len: int, embed_dim: int):
54
+ super().__init__()
55
+ self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
56
+
57
+ def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
58
+ seq_len = x.size(1)
59
+ if pos_offset + seq_len > self.pos_emb.size(0):
60
+ raise ValueError("Sequence length exceeds MAX_SEQ_LEN defined in position embedding.")
61
+
62
+ pos = self.pos_emb[pos_offset : pos_offset + seq_len]
63
+ return x + pos.unsqueeze(0)
64
+
65
+
66
+ # -------------------------------
67
+ # MultiHeadAttention (MHA)
68
+ # -------------------------------
69
+ class MultiHeadAttention(nn.Module):
70
+ def __init__(self):
71
+ super().__init__()
72
+ self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
73
+ self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
74
+ self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
75
+ self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
76
+ self.scale = HEAD_DIM ** -0.5
77
+
78
+ def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
79
+ B, T, D = x.shape
80
+
81
+ q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
82
+ k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
83
+ v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
84
+
85
+ # 1. KV-кеш и определение смещения
86
+ pos_offset = 0
87
+ seqlen_k_new = k.size(2)
88
+ if past_kv is not None:
89
+ past_k, past_v = past_kv
90
+ k = torch.cat([past_k, k], dim=2)
91
+ v = torch.cat([past_v, v], dim=2)
92
+ pos_offset = past_k.size(2)
93
+
94
+ seqlen_k = k.size(2) # Общая длина K
95
+ new_kv = (k, v)
96
+
97
+ # 2. Расчет внимания
98
+ attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
99
+
100
+ # 3. КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ МАСКИРОВАНИЯ (Causal Mask)
101
+ if T == seqlen_k_new and seqlen_k > 0:
102
+ # Создаем маску T (query length) x seqlen_k (key length)
103
+ mask = torch.full((T, seqlen_k),
104
+ float("-inf"),
105
+ device=x.device,
106
+ dtype=attn.dtype)
107
+
108
+ # Разрешаем всем новым токенам видеть все старые токены (past_kv)
109
+ mask[:, :pos_offset] = 0.0
110
+
111
+ # Применяем треугольную маску для текущих T токенов
112
+ current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
113
+ mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
114
+
115
+ # Применяем маску к весам внимания
116
+ attn = attn + mask[None, None, :, :]
117
+
118
+ # 4. Выход
119
+ attn = F.softmax(attn, dim=-1)
120
+ out = torch.matmul(attn, v)
121
+ out = out.transpose(1, 2).contiguous().view(B, T, D)
122
+ out = self.out_proj(out)
123
+
124
+ return out, new_kv
125
+
126
+
127
+ # -------------------------------
128
+ # FeedForward (GELU, GPT-style)
129
+ # -------------------------------
130
+ class FeedForward(nn.Module):
131
+ def __init__(self):
132
+ super().__init__()
133
+ self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
134
+ self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
135
+
136
+ def forward(self, x):
137
+ return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
138
+
139
+
140
+ # -------------------------------
141
+ # Transformer Block (Post-Norm, GPT-style)
142
+ # -------------------------------
143
+ class TransformerBlock(nn.Module):
144
+ def __init__(self):
145
+ super().__init__()
146
+ self.attn = MultiHeadAttention()
147
+ self.ffn = FeedForward()
148
+ # ИСПРАВЛЕНИЕ: Добавлен стандартный eps
149
+ self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
150
+ self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
151
+
152
+ def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
153
+ # Post-Normalization (GPT Style)
154
+ attn_out, new_kv = self.attn(self.norm1(x), past_kv)
155
+ x = x + attn_out
156
+ x = x + self.ffn(self.norm2(x))
157
+ return x, new_kv
158
+
159
+
160
+ # -------------------------------
161
+ # Главная модель GPTPyTorch (L=6, H=8)
162
+ # -------------------------------
163
+ class GPTPyTorch(nn.Module):
164
+ def __init__(self):
165
+ super().__init__()
166
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
167
+ self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
168
+ self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
169
+ self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
170
+ self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
171
+
172
+ signature = "Konstantin V Gbabko . original author © 2025"
173
+ bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
174
+ self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
175
+ self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
176
+
177
+ self.lm_head.weight = self.token_emb.weight
178
+ self.apply(self._init_weights)
179
+
180
+ def _init_weights(self, module):
181
+ if isinstance(module, nn.Linear):
182
+ # ИСПРАВЛЕНИЕ: Инициализация, масштабированная по глубине сети.
183
+ std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
184
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
185
+ elif isinstance(module, nn.Embedding):
186
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
187
+ elif isinstance(module, nn.LayerNorm):
188
+ nn.init.zeros_(module.bias)
189
+ nn.init.ones_(module.weight)
190
+
191
+ def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
192
+ B, T = input_ids.shape
193
+ x = self.token_emb(input_ids)
194
+
195
+ pos_offset = 0
196
+ if past_kv is not None and past_kv[0] is not None:
197
+ pos_offset = past_kv[0][0].size(2)
198
+
199
+ x = self.pos_emb(x, pos_offset=pos_offset)
200
+
201
+ # Инициализация нового кеша
202
+ new_kv_cache = [] if past_kv is not None or T > 1 else None
203
+ current_past = past_kv
204
+
205
+ for i, block in enumerate(self.blocks):
206
+ layer_past = current_past[i] if (current_past and i < len(current_past)) else None
207
+ x, layer_kv = block(x, layer_past)
208
+
209
+ if new_kv_cache is not None:
210
+ new_kv_cache.append(layer_kv)
211
+
212
+ x = self.ln_f(x)
213
+ logits = self.lm_head(x)
214
+ return logits, new_kv_cache
215
+
216
+ @torch.no_grad()
217
+ def generate(
218
+ self,
219
+ input_ids: torch.Tensor,
220
+ max_new_tokens: int = 100,
221
+ temperature: float = 0.8,
222
+ top_p: float = 0.95,
223
+ repetition_penalty: float = 1.0,
224
+ do_sample: bool = True,
225
+ eos_token_id: int = 50256
226
+ ) -> torch.Tensor:
227
+ kv_cache = [None] * NUM_LAYERS
228
+ current_ids = input_ids.clone()
229
+
230
+ for step in range(max_new_tokens):
231
+ if step == 0:
232
+ input_for_model = current_ids
233
+ else:
234
+ input_for_model = current_ids[:, -1].unsqueeze(-1)
235
+
236
+ logits, kv_cache = self(input_for_model, kv_cache)
237
+ next_token_logits = logits[:, -1, :]
238
+
239
+ if temperature > 0:
240
+ next_token_logits = next_token_logits / temperature
241
+
242
+ # Repetition Penalty (логика сохранена)
243
+ if repetition_penalty != 1.0:
244
+ for i in range(current_ids.shape[0]):
245
+ unique_tokens = torch.unique(current_ids[i]).tolist()
246
+ for token_id in unique_tokens:
247
+ score = next_token_logits[i, token_id]
248
+ if score < 0:
249
+ next_token_logits[i, token_id] = score * repetition_penalty
250
+ else:
251
+ next_token_logits[i, token_id] = score / repetition_penalty
252
+
253
+ # Top-P сэмплирование (логика сохранена)
254
+ if do_sample and top_p < 1.0:
255
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
256
+ cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
257
+ sorted_indices_to_remove = cumulative_probs > top_p
258
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
259
+ sorted_indices_to_remove[:, 0] = False
260
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
261
+ next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))
262
+
263
+ # Сэмплирование
264
+ if do_sample and temperature > 0:
265
+ probs = torch.softmax(next_token_logits, dim=-1)
266
+ if torch.isnan(probs).any() or torch.isinf(probs).any():
267
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
268
+ else:
269
+ next_token = torch.multinomial(probs, num_samples=1)
270
+ else:
271
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
272
+
273
+ if next_token.item() == eos_token_id:
274
+ break
275
+
276
+ current_ids = torch.cat([current_ids, next_token], dim=1)
277
+
278
+ return current_ids
279
+
280
+
281
+ if __name__ == "__main__":
282
+ os.makedirs("models", exist_ok=True)
283
+
284
+ model = GPTPyTorch().to(device)
285
+ model.eval()
286
+
287
+ total_params = sum(p.numel() for p in model.parameters())
288
+ print(f"Device: {device}")
289
+ print(f"Total parameters: {total_params / 1e6:.2f}M") # ~124.44M
290
+
291
+ # 1. Проверка первого прохода (T=50)
292
+ input_ids_T50 = torch.randint(0, VOCAB_SIZE, (1, 50), device=device)
293
+ with torch.no_grad():
294
+ logits_50, kv_cache_50 = model(input_ids_T50)
295
+ expected_k_shape = (1, NUM_HEADS, 50, HEAD_DIM)
296
+ # Проверяем, что размерность головы HEAD_DIM (96) соблюдена
297
+ assert kv_cache_50[0][0].shape == expected_k_shape
298
+ print(f"Initial logits shape: {logits_50.shape}")
299
+ print(f"Initial KV-cache seqlen: {kv_cache_50[0][0].size(2)} (Head Dim: {kv_cache_50[0][0].size(3)})")
300
+
301
+ # 2. Проверка инкрементального прохода (T=1)
302
+ input_ids_T1 = torch.randint(0, VOCAB_SIZE, (1, 1), device=device)
303
+ with torch.no_grad():
304
+ logits_51, kv_cache_51 = model(input_ids_T1, past_kv=kv_cache_50)
305
+
306
+ # Проверка длины кеша: 50 + 1 = 51
307
+ assert kv_cache_51[0][0].size(2) == 51
308
+ print(f"Incremental logits shape: {logits_51.shape}")
309
+ print(f"Incremental KV-cache seqlen: {kv_cache_51[0][0].size(2)}")
310
+
311
+ # 3. Проверка функции generate
312
+ generated = model.generate(input_ids_T50, max_new_tokens=5, temperature=0.8, top_p=0.9)
313
+ print(f"Generated sequence length (50 + 5): {generated.shape[1]}")
314
+
315
+ save_path = "models/JiRack_GPT_L6_H8_PostNorm_fixed.pt"
316
+ torch.save(model.state_dict(), save_path)
317
+ print(f"Model successfully saved to {save_path}")
source/JiRack_H8_L8_V50257_D768_MSL8192_FF768x4.py ADDED
@@ -0,0 +1,317 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2025 CMS Manhattan
2
+ # All rights reserved.
3
+ # Author: Konstantin Vladimirovich Grabko
4
+ # Email: grabko@cmsmanhattan.com
5
+ # Phone: +1(516)777-0945
6
+ #
7
+ # This program is free software: you can redistribute it and/or modify
8
+ # it under the terms of the GNU General Public License as published by
9
+ # the Free Software Foundation, version 3 of the License.
10
+ #
11
+ # This program is distributed in the hope that it will be useful,
12
+ # but WITHOUT ANY WARRANTY; without even the implied warranty of
13
+ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14
+ # GNU General Public License for more details.
15
+ #
16
+ # You should have received a copy of the GNU General Public License
17
+ # along with this program. If not, see <https://www.gnu.org/licenses/>.
18
+ #
19
+ # Additional terms:
20
+ # Any commercial use or distribution of this software or derivative works
21
+ # requires explicit written permission from the copyright holder.
22
+
23
+ import os
24
+ import torch
25
+ import torch.nn as nn
26
+ import torch.nn.functional as F
27
+ from typing import Optional, Tuple, List
28
+ import math
29
+
30
+ # ========================================
31
+ # Model Configuration (L=8, H=8, D=768)
32
+ # ========================================
33
+ VOCAB_SIZE = 50257
34
+ MODEL_DIM = 768
35
+ NUM_HEADS = 8
36
+ NUM_LAYERS = 8 # Set to 8 layers
37
+ MAX_SEQ_LEN = 8192
38
+ FFN_HIDDEN_DIM = 4 * MODEL_DIM
39
+ HEAD_DIM = MODEL_DIM // NUM_HEADS # 768 / 8 = 96
40
+
41
+ # ИСПРАВЛЕНИЕ: ROCm/HIP-совместимая проверка устройства
42
+ if torch.cuda.is_available():
43
+ device = torch.device("cuda")
44
+ elif hasattr(torch, 'hip') and torch.hip.is_available():
45
+ device = torch.device("cuda")
46
+ else:
47
+ device = torch.device("cpu")
48
+
49
+ # -------------------------------
50
+ # Learned Positional Embedding
51
+ # -------------------------------
52
+ class LearnedPositionalEmbedding(nn.Module):
53
+ def __init__(self, max_seq_len: int, embed_dim: int):
54
+ super().__init__()
55
+ self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
56
+
57
+ def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
58
+ seq_len = x.size(1)
59
+ if pos_offset + seq_len > self.pos_emb.size(0):
60
+ raise ValueError("Sequence length exceeds MAX_SEQ_LEN defined in position embedding.")
61
+
62
+ pos = self.pos_emb[pos_offset : pos_offset + seq_len]
63
+ return x + pos.unsqueeze(0)
64
+
65
+
66
+ # -------------------------------
67
+ # MultiHeadAttention (MHA)
68
+ # -------------------------------
69
+ class MultiHeadAttention(nn.Module):
70
+ def __init__(self):
71
+ super().__init__()
72
+ self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
73
+ self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
74
+ self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
75
+ self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
76
+ self.scale = HEAD_DIM ** -0.5
77
+
78
+ def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
79
+ B, T, D = x.shape
80
+
81
+ q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
82
+ k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
83
+ v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
84
+
85
+ # 1. KV-кеш и определение смещения
86
+ pos_offset = 0
87
+ seqlen_k_new = k.size(2)
88
+ if past_kv is not None:
89
+ past_k, past_v = past_kv
90
+ k = torch.cat([past_k, k], dim=2)
91
+ v = torch.cat([past_v, v], dim=2)
92
+ pos_offset = past_k.size(2)
93
+
94
+ seqlen_k = k.size(2) # Общая длина K
95
+ new_kv = (k, v)
96
+
97
+ # 2. Расчет внимания
98
+ attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
99
+
100
+ # 3. КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ МАСКИРОВАНИЯ (Causal Mask)
101
+ if T == seqlen_k_new and seqlen_k > 0:
102
+ # Создаем маску T (query length) x seqlen_k (key length)
103
+ mask = torch.full((T, seqlen_k),
104
+ float("-inf"),
105
+ device=x.device,
106
+ dtype=attn.dtype)
107
+
108
+ # Разрешаем всем новым токенам видеть все старые токены (past_kv)
109
+ mask[:, :pos_offset] = 0.0
110
+
111
+ # Применяем треугольную маску для текущих T токенов
112
+ current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
113
+ mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
114
+
115
+ # Применяем маску к весам внимания
116
+ attn = attn + mask[None, None, :, :]
117
+
118
+ # 4. Выход
119
+ attn = F.softmax(attn, dim=-1)
120
+ out = torch.matmul(attn, v)
121
+ out = out.transpose(1, 2).contiguous().view(B, T, D)
122
+ out = self.out_proj(out)
123
+
124
+ return out, new_kv
125
+
126
+
127
+ # -------------------------------
128
+ # FeedForward (GELU, GPT-style)
129
+ # -------------------------------
130
+ class FeedForward(nn.Module):
131
+ def __init__(self):
132
+ super().__init__()
133
+ self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
134
+ self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
135
+
136
+ def forward(self, x):
137
+ return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
138
+
139
+
140
+ # -------------------------------
141
+ # Transformer Block (Post-Norm, GPT-style)
142
+ # -------------------------------
143
+ class TransformerBlock(nn.Module):
144
+ def __init__(self):
145
+ super().__init__()
146
+ self.attn = MultiHeadAttention()
147
+ self.ffn = FeedForward()
148
+ # ИСПРАВЛЕНИЕ: Добавлен стандартный eps
149
+ self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
150
+ self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
151
+
152
+ def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
153
+ # Post-Normalization (GPT Style)
154
+ attn_out, new_kv = self.attn(self.norm1(x), past_kv)
155
+ x = x + attn_out
156
+ x = x + self.ffn(self.norm2(x))
157
+ return x, new_kv
158
+
159
+
160
+ # -------------------------------
161
+ # Главная модель GPTPyTorch (L=8, H=8)
162
+ # -------------------------------
163
+ class GPTPyTorch(nn.Module):
164
+ def __init__(self):
165
+ super().__init__()
166
+ self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
167
+ self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
168
+ self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
169
+ self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
170
+ self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
171
+
172
+ signature = "Konstantin V Gbabko . original author © 2025"
173
+ bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
174
+ self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
175
+ self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
176
+
177
+ self.lm_head.weight = self.token_emb.weight
178
+ self.apply(self._init_weights)
179
+
180
+ def _init_weights(self, module):
181
+ if isinstance(module, nn.Linear):
182
+ # ИСПРАВЛЕНИЕ: Инициализация, масштабированная по глубине сети.
183
+ std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
184
+ torch.nn.init.normal_(module.weight, mean=0.0, std=std)
185
+ elif isinstance(module, nn.Embedding):
186
+ torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
187
+ elif isinstance(module, nn.LayerNorm):
188
+ nn.init.zeros_(module.bias)
189
+ nn.init.ones_(module.weight)
190
+
191
+ def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
192
+ B, T = input_ids.shape
193
+ x = self.token_emb(input_ids)
194
+
195
+ pos_offset = 0
196
+ if past_kv is not None and past_kv[0] is not None:
197
+ pos_offset = past_kv[0][0].size(2)
198
+
199
+ x = self.pos_emb(x, pos_offset=pos_offset)
200
+
201
+ # Инициализация нового кеша
202
+ new_kv_cache = [] if past_kv is not None or T > 1 else None
203
+ current_past = past_kv
204
+
205
+ for i, block in enumerate(self.blocks):
206
+ layer_past = current_past[i] if (current_past and i < len(current_past)) else None
207
+ x, layer_kv = block(x, layer_past)
208
+
209
+ if new_kv_cache is not None:
210
+ new_kv_cache.append(layer_kv)
211
+
212
+ x = self.ln_f(x)
213
+ logits = self.lm_head(x)
214
+ return logits, new_kv_cache
215
+
216
+ @torch.no_grad()
217
+ def generate(
218
+ self,
219
+ input_ids: torch.Tensor,
220
+ max_new_tokens: int = 100,
221
+ temperature: float = 0.8,
222
+ top_p: float = 0.95,
223
+ repetition_penalty: float = 1.0,
224
+ do_sample: bool = True,
225
+ eos_token_id: int = 50256
226
+ ) -> torch.Tensor:
227
+ kv_cache = [None] * NUM_LAYERS
228
+ current_ids = input_ids.clone()
229
+
230
+ for step in range(max_new_tokens):
231
+ if step == 0:
232
+ input_for_model = current_ids
233
+ else:
234
+ input_for_model = current_ids[:, -1].unsqueeze(-1)
235
+
236
+ logits, kv_cache = self(input_for_model, kv_cache)
237
+ next_token_logits = logits[:, -1, :]
238
+
239
+ if temperature > 0:
240
+ next_token_logits = next_token_logits / temperature
241
+
242
+ # Repetition Penalty (логика сохранена)
243
+ if repetition_penalty != 1.0:
244
+ for i in range(current_ids.shape[0]):
245
+ unique_tokens = torch.unique(current_ids[i]).tolist()
246
+ for token_id in unique_tokens:
247
+ score = next_token_logits[i, token_id]
248
+ if score < 0:
249
+ next_token_logits[i, token_id] = score * repetition_penalty
250
+ else:
251
+ next_token_logits[i, token_id] = score / repetition_penalty
252
+
253
+ # Top-P сэмплирование (логика сохранена)
254
+ if do_sample and top_p < 1.0:
255
+ sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
256
+ cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
257
+ sorted_indices_to_remove = cumulative_probs > top_p
258
+ sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
259
+ sorted_indices_to_remove[:, 0] = False
260
+ indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
261
+ next_token_logits = next_token_logits.masked_fill(indices_to_remove, float('-inf'))
262
+
263
+ # Сэмплирование
264
+ if do_sample and temperature > 0:
265
+ probs = torch.softmax(next_token_logits, dim=-1)
266
+ if torch.isnan(probs).any() or torch.isinf(probs).any():
267
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
268
+ else:
269
+ next_token = torch.multinomial(probs, num_samples=1)
270
+ else:
271
+ next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
272
+
273
+ if next_token.item() == eos_token_id:
274
+ break
275
+
276
+ current_ids = torch.cat([current_ids, next_token], dim=1)
277
+
278
+ return current_ids
279
+
280
+
281
+ if __name__ == "__main__":
282
+ os.makedirs("models", exist_ok=True)
283
+
284
+ model = GPTPyTorch().to(device)
285
+ model.eval()
286
+
287
+ total_params = sum(p.numel() for p in model.parameters())
288
+ print(f"Device: {device}")
289
+ print(f"Total parameters: {total_params / 1e6:.2f}M") # ~157.14M
290
+
291
+ # 1. Проверка первого прохода (T=50)
292
+ input_ids_T50 = torch.randint(0, VOCAB_SIZE, (1, 50), device=device)
293
+ with torch.no_grad():
294
+ logits_50, kv_cache_50 = model(input_ids_T50)
295
+ expected_k_shape = (1, NUM_HEADS, 50, HEAD_DIM)
296
+ # Проверяем, что размерность головы HEAD_DIM (96) соблюдена
297
+ assert kv_cache_50[0][0].shape == expected_k_shape
298
+ print(f"Initial logits shape: {logits_50.shape}")
299
+ print(f"Initial KV-cache seqlen: {kv_cache_50[0][0].size(2)} (Head Dim: {kv_cache_50[0][0].size(3)})")
300
+
301
+ # 2. Проверка инкрементального прохода (T=1)
302
+ input_ids_T1 = torch.randint(0, VOCAB_SIZE, (1, 1), device=device)
303
+ with torch.no_grad():
304
+ logits_51, kv_cache_51 = model(input_ids_T1, past_kv=kv_cache_50)
305
+
306
+ # Проверка длины кеша: 50 + 1 = 51
307
+ assert kv_cache_51[0][0].size(2) == 51
308
+ print(f"Incremental logits shape: {logits_51.shape}")
309
+ print(f"Incremental KV-cache seqlen: {kv_cache_51[0][0].size(2)}")
310
+
311
+ # 3. Проверка функции generate
312
+ generated = model.generate(input_ids_T50, max_new_tokens=5, temperature=0.8, top_p=0.9)
313
+ print(f"Generated sequence length (50 + 5): {generated.shape[1]}")
314
+
315
+ save_path = "models/JiRack_GPT_L8_H8_PostNorm_fixed.pt"
316
+ torch.save(model.state_dict(), save_path)
317
+ print(f"Model successfully saved to {save_path}")