Upload 16 files

source_jit/JiRack_H12_L12_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,266 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+
+# ========================================
+# Model Configuration (GPT-2 Base Style)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_HEADS = 12
+NUM_LAYERS = 12
+MAX_SEQ_LEN = 8192
+FFN_HIDDEN_DIM = 4 * MODEL_DIM
+HEAD_DIM = MODEL_DIM // NUM_HEADS
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+# -------------------------------
+# Learned Positional Embedding
+# -------------------------------
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, embed_dim: int):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
+
+    def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
+        seq_len = x.size(1)
+        # Проверка на длину убрана для JIT-совместимости (TracerWarning)
+        
+        pos = self.pos_emb[pos_offset : pos_offset + seq_len]
+        return x + pos.unsqueeze(0)
+
+# -------------------------------
+# MultiHeadAttention 
+# -------------------------------
+class MultiHeadAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.scale = HEAD_DIM ** -0.5
+
+    def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, D = x.shape
+        
+        q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        # Обработка KV-кэша
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            pos_offset = past_k.size(2)
+            new_kv = (k, v)
+        else:
+            new_kv = None 
+
+        seqlen_k = k.size(2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # === ОКОНЧАТЕЛЬНОЕ ИСПРАВЛЕНИЕ: Безусловное маскирование для устранения TracerWarning ===
+        # Проверяем T > 0, но не в условном Python 'if'
+        mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
+        
+        # Маска для текущего блока (T x T)
+        current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
+        
+        # Маска для past_kv (T x pos_offset)
+        mask[:, :pos_offset] = 0.0
+        
+        # Применение текущей причинной маски к части K, соответствующей текущему входу
+        mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
+        attn = attn + mask[None, None, :, :]
+
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.out_proj(out)
+
+        return out, new_kv 
+# -------------------------------
+# FeedForward 
+# -------------------------------
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
+        self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
+
+# -------------------------------
+# Transformer Block 
+# -------------------------------
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = MultiHeadAttention()
+        self.ffn = FeedForward()
+        self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5) 
+        self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+    def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, new_kv = self.attn(self.norm1(x), past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        
+        return x, new_kv
+# -------------------------------
+# Главная модель GPTPyTorch
+# -------------------------------
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+        
+        signature = "Konstantin V Gbabko .  original author © 2025"
+        bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
+        self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
+        self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            std = 0.02 / (2 * NUM_LAYERS)**0.5 if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None and isinstance(past_kv[0], tuple):
+            pos_offset = past_kv[0][0].size(2)
+            
+        x = self.pos_emb(x, pos_offset=pos_offset)
+
+        new_kv_cache = [] if past_kv is not None else None
+        current_past = past_kv
+
+        for i, block in enumerate(self.blocks):
+            layer_past = current_past[i] if (current_past and i < len(current_past)) else None
+            
+            x, layer_kv = block(x, layer_past)
+
+            if past_kv is not None:
+                new_kv_cache.append(layer_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        if past_kv is None:
+            return logits
+        else:
+            return logits, new_kv_cache
+
+# -------------------------------
+# Обертка для JIT-трассировки (гарантирует только Tensor)
+# -------------------------------
+class GPTPyTorchNoCache(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        
+    def forward(self, input_ids):
+        return self.model(input_ids, None)
+
+
+# -------------------------------
+# ОСНОВНОЙ БЛОК: JIT-КОНВЕРТАЦИЯ
+# -------------------------------
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+
+    # =========================================================================
+    # 1. ПОДГОТОВКА МОДЕЛИ И ПРОВЕРКА
+    # =========================================================================
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Device: {device}")
+    print(f"Total parameters: {total_params / 1e6:.2f}M")
+
+    TRAIN_SEQ_LEN = 256
+    
+    # Создаем фиктивный входной тензор
+    dummy_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device)
+    
+    # Проверяем обычный проход
+    with torch.no_grad():
+        logits_test = model(dummy_input, None)
+        print(f"Test logits shape: {logits_test.shape}")
+        
+    # =========================================================================
+    # 2. JIT-ТРАССИРОВКА И СОХРАНЕНИЕ
+    # =========================================================================
+    print(f"\nTracing model for JIT export (input sequence length: {TRAIN_SEQ_LEN})...")
+    
+    # Используем обертку для тр��ссировки
+    model_no_cache = GPTPyTorchNoCache(model).to(device)
+
+    try:
+        # Трассируем обертку, которая принимает только input_ids
+        traced_script_module = torch.jit.trace(model_no_cache, dummy_input, strict=False)
+        JIT_SAVE_PATH = "models/JiRack_H12_L12_V50257_D768_MSL8192_FF768x4.script.pt"
+        traced_script_module.save(JIT_SAVE_PATH)
+        print(f"✅ Success! Model saved as TorchScript (JIT) to: {JIT_SAVE_PATH}")
+        print("Now you can run your training script.")
+    except Exception as e:
+        print(f"🚨 ERROR during JIT tracing: {e}")
+        print("Model may contain operations incompatible with torch.jit.trace.")
+        
+    # Сохраняем оригинальную модель (на всякий случай)
+    ORIGINAL_SAVE_PATH = "models/JiRack_H12_L12_V50257_D768_MSL8192_FF768x4.pt"
+    torch.save(model.state_dict(), ORIGINAL_SAVE_PATH)
+    print(f"\nOriginal state_dict saved to {ORIGINAL_SAVE_PATH}")

source_jit/JiRack_H12_L18_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,248 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+from pathlib import Path
+
+# ========================================
+# Model Configuration (L=18, D=768)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_HEADS = 12
+NUM_LAYERS = 18       # Increased depth
+MAX_SEQ_LEN = 8192
+FFN_HIDDEN_DIM = 4 * MODEL_DIM
+HEAD_DIM = MODEL_DIM // NUM_HEADS
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+# --- Learned Positional Embedding ---
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, embed_dim: int):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
+
+    def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
+        seq_len = x.size(1)
+        pos = self.pos_emb[pos_offset : pos_offset + seq_len]
+        return x + pos.unsqueeze(0)
+
+
+# --- MultiHeadAttention ---
+class MultiHeadAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.scale = HEAD_DIM ** -0.5
+
+    def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, D = x.shape
+        
+        q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        pos_offset = 0
+        new_kv = None
+        if past_kv is not None and past_kv[0] is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            pos_offset = past_k.size(2)
+            new_kv = (k, v) 
+
+        seqlen_k = k.size(2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # === ФИНАЛЬНОЕ ИСПРАВЛЕНИЕ (УДАЛЕНИЕ T > 0) ===
+        # Удалено 'if T > 0:', чтобы избежать TracerWarning при трассировке.
+        # Код маскирования теперь всегда выполняется (для T >= 1).
+        mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
+        current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
+        
+        mask[:, :pos_offset] = 0.0
+        mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
+        
+        attn = attn + mask[None, None, :, :]
+        # ===============================================
+
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.out_proj(out)
+
+        return out, new_kv
+
+
+# --- FeedForward (Без изменений) ---
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
+        self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
+
+
+# --- Transformer Block (Без изменений) ---
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = MultiHeadAttention()
+        self.ffn = FeedForward()
+        self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5) 
+        self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+    def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, new_kv = self.attn(self.norm1(x), past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+# -------------------------------
+# Главная модель GPTPyTorch
+# -------------------------------
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+        
+        signature = "Konstantin V Gbabko .  original author © 2025"
+        bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
+        self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
+        self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None:
+            pos_offset = past_kv[0][0].size(2)
+            
+        x = self.pos_emb(x, pos_offset=pos_offset)
+
+        new_kv_cache = [] if past_kv is not None else None
+        current_past = past_kv
+
+        for i, block in enumerate(self.blocks):
+            layer_past = current_past[i] if (current_past and i < len(current_past)) else None
+            x, layer_kv = block(x, layer_past)
+            
+            if new_kv_cache is not None:
+                new_kv_cache.append(layer_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        
+        if past_kv is None:
+            return logits 
+        else:
+            return logits, new_kv_cache
+
+# -------------------------------
+# Обертка для JIT-трассировки (гарантирует только Tensor)
+# -------------------------------
+class GPTPyTorchNoCache(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        
+    def forward(self, input_ids):
+        return self.model(input_ids, None)
+
+# =========================================================================
+# ОСНОВНОЙ БЛОК: JIT-КОНВЕРТАЦИЯ
+# =========================================================================
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+    
+    TRAIN_SEQ_LEN = 256
+    JIT_SAVE_PATH = Path("models/JiRack_H12_L18_V50257_D768_MSL8192_FF768x4.script.pt")
+    
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Device: {device}")
+    print(f"Total parameters: {total_params / 1e6:.2f}M")
+    
+    # 1. Проверка первого прохода
+    dummy_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device)
+    with torch.no_grad():
+        logits_test = model(dummy_input, None) 
+        print(f"Test logits shape: {logits_test.shape}")
+        
+    # 2. JIT-ТРАССИРОВКА И СОХРАНЕНИЕ
+    print(f"\nTracing model for JIT export (input sequence length: {TRAIN_SEQ_LEN})...")
+    
+    model_no_cache = GPTPyTorchNoCache(model).to(device)
+
+    try:
+        traced_script_module = torch.jit.trace(model_no_cache, dummy_input, strict=False)
+        
+        traced_script_module.save(JIT_SAVE_PATH)
+        print(f"✅ Success! Model saved as TorchScript (JIT) to: {JIT_SAVE_PATH}")
+        print("Now you can run your training script.")
+        
+    except Exception as e:
+        print(f"🚨 ERROR during JIT tracing: {e}")
+        print("Model may contain operations incompatible with torch.jit.trace.")
+        
+    # Сохраняем оригинальную модель (на всякий случай)
+    ORIGINAL_SAVE_PATH = "models/JiRack_H12_L18_V50257_D768_MSL8192_FF768x4.pt"
+    torch.save(model.state_dict(), ORIGINAL_SAVE_PATH)
+    print(f"\nOriginal state_dict saved to {ORIGINAL_SAVE_PATH}")

source_jit/JiRack_H12_L24_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,264 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+from pathlib import Path
+
+# ========================================
+# Model Configuration (GPT-2 Large Style)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_HEADS = 12
+NUM_LAYERS = 24       # Increased depth (GPT-2 Large equivalent)
+MAX_SEQ_LEN = 8192
+FFN_HIDDEN_DIM = 4 * MODEL_DIM
+HEAD_DIM = MODEL_DIM // NUM_HEADS
+
+# ROCm/HIP-совместимая проверка устройства
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+# --- Learned Positional Embedding (Без изменений) ---
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, embed_dim: int):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
+
+    def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
+        seq_len = x.size(1)
+        # Убрана Python-проверка для JIT-совместимости (если T > MAX_SEQ_LEN)
+        pos = self.pos_emb[pos_offset : pos_offset + seq_len]
+        return x + pos.unsqueeze(0)
+
+
+# --- MultiHeadAttention (MHA) (Исправлено для JIT) ---
+class MultiHeadAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.scale = HEAD_DIM ** -0.5
+
+    def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, D = x.shape
+        
+        q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        pos_offset = 0
+        new_kv = None
+        
+        # --- KV-кэш логика ---
+        if past_kv is not None and past_kv[0] is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            pos_offset = past_k.size(2)
+            new_kv = (k, v) 
+        elif past_kv is not None:
+            # Случай, когда past_kv передан, но пуст (например, [None]*24)
+            new_kv = (k, v)
+        
+        seqlen_k = k.size(2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # === ИСПРАВЛЕНИЕ TRACERWARNING: Убрано if/and, выполняется безусловно для T >= 1 ===
+        # Поскольку T=256 при трассировке, этот путь всегда записывается.
+        mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
+        current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
+        
+        mask[:, :pos_offset] = 0.0
+        mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
+        
+        attn = attn + mask[None, None, :, :]
+        # ===================================================================================
+
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.out_proj(out)
+
+        return out, new_kv
+
+
+# --- FeedForward (Без изменений) ---
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
+        self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
+
+
+# --- Transformer Block (Без изменений) ---
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = MultiHeadAttention()
+        self.ffn = FeedForward()
+        self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5) 
+        self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+    def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, new_kv = self.attn(self.norm1(x), past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+# -------------------------------
+# Главная модель GPTPyTorch (24 слоя)
+# -------------------------------
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+        
+        signature = "Konstantin V Gbabko .  original author © 2025"
+        bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
+        self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
+        self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # Инициализация, масштабированная по глубине сети (L=24)
+            std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    # Метод forward для обучения и инференса с кешем
+    def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None:
+            pos_offset = past_kv[0][0].size(2)
+            
+        x = self.pos_emb(x, pos_offset=pos_offset)
+
+        new_kv_cache = [] if past_kv is not None else None
+        current_past = past_kv
+
+        for i, block in enumerate(self.blocks):
+            layer_past = current_past[i] if (current_past and i < len(current_past)) else None
+            x, layer_kv = block(x, layer_past)
+            
+            if new_kv_cache is not None:
+                new_kv_cache.append(layer_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        
+        # === КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: Статический вывод для JIT-трассировки ===
+        if past_kv is None:
+            # Путь обучения: JIT может трассировать только Tensor
+            return logits 
+        else:
+            # Путь инференса с кэшем: возвращаем Tensor и кеш
+            return logits, new_kv_cache
+
+# -------------------------------
+# Обертка для JIT-трассировки (гарантирует только Tensor)
+# -------------------------------
+class GPTPyTorchNoCache(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        
+    def forward(self, input_ids):
+        # Вызов с None гарантирует, что основной forward вернет только logits (Tensor).
+        return self.model(input_ids, None)
+
+# =========================================================================
+# ОСНОВНОЙ БЛОК: JIT-КОНВЕРТАЦИЯ
+# =========================================================================
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+    
+    TRAIN_SEQ_LEN = 256
+    # Обновленное имя файла для отражения L=24
+    JIT_SAVE_PATH = Path("models/JiRack_H12_L24_V50257_D768_MSL8192_FF768x4.script.pt")
+    
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Device: {device}")
+    print(f"Total parameters: {total_params / 1e6:.2f}M")
+    
+    # 1. Проверка первого прохода
+    dummy_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device)
+    with torch.no_grad():
+        # Тестируем путь обучения (past_kv=None)
+        logits_test = model(dummy_input, None) 
+        print(f"Test logits shape: {logits_test.shape}")
+        
+    # 2. JIT-ТРАССИРОВКА И СОХРАНЕНИЕ
+    print(f"\nTracing model for JIT export (input sequence length: {TRAIN_SEQ_LEN})...")
+    
+    model_no_cache = GPTPyTorchNoCache(model).to(device)
+
+    try:
+        # Трассируем обертку, которая принимает только input_ids, 
+        # что гарантирует один Tensor на выходе.
+        traced_script_module = torch.jit.trace(model_no_cache, dummy_input, strict=False)
+        
+        traced_script_module.save(JIT_SAVE_PATH)
+        print(f"✅ Success! Model saved as TorchScript (JIT) to: {JIT_SAVE_PATH}")
+        print("Now you can run your training script to fine-tune this model.")
+        
+    except Exception as e:
+        print(f"🚨 ERROR during JIT tracing: {e}")
+        print("Model may contain operations incompatible with torch.jit.trace.")
+        
+    # Сохраняем оригинальную модель (на всякий случай)
+    ORIGINAL_SAVE_PATH = "models/JiRack_H12_L24_V50257_D768_MSL8192_FF768x4.pt"
+    torch.save(model.state_dict(), ORIGINAL_SAVE_PATH)
+    print(f"\nOriginal state_dict saved to {ORIGINAL_SAVE_PATH}")

source_jit/JiRack_H12_L32_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,260 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+from pathlib import Path
+
+# ========================================
+# Model Configuration (GPT-2 XL Style)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_HEADS = 12
+NUM_LAYERS = 32       # Increased depth (GPT-2 XL equivalent)
+MAX_SEQ_LEN = 8192
+FFN_HIDDEN_DIM = 4 * MODEL_DIM
+HEAD_DIM = MODEL_DIM // NUM_HEADS
+
+# ROCm/HIP-совместимая проверка устройства
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+# --- Learned Positional Embedding (Исправлено для JIT) ---
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, embed_dim: int):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
+
+    def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
+        seq_len = x.size(1)
+        # ИСПРАВЛЕНИЕ: Удалена Python-проверка, несовместимая с JIT
+        pos = self.pos_emb[pos_offset : pos_offset + seq_len]
+        return x + pos.unsqueeze(0)
+
+
+# --- MultiHeadAttention (MHA) (Исправлено для JIT) ---
+class MultiHeadAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.scale = HEAD_DIM ** -0.5
+
+    def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, D = x.shape
+        
+        q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        pos_offset = 0
+        new_kv = None
+        
+        if past_kv is not None and past_kv[0] is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            pos_offset = past_k.size(2)
+            new_kv = (k, v)
+        elif past_kv is not None:
+            new_kv = (k, v) 
+        
+        seqlen_k = k.size(2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # ИСПРАВЛЕНИЕ: Удалена динамическая Python-проверка 'if T == seqlen_k_new and seqlen_k > 0:'
+        # Маскирование выполняется безусловно для JIT-совместимости.
+        mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
+        current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
+        
+        mask[:, :pos_offset] = 0.0
+        mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
+        
+        attn = attn + mask[None, None, :, :]
+
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.out_proj(out)
+
+        return out, new_kv
+
+
+# --- FeedForward (Без изменений) ---
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
+        self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
+
+
+# --- Transformer Block (Без изменений) ---
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = MultiHeadAttention()
+        self.ffn = FeedForward()
+        self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5) 
+        self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+    def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, new_kv = self.attn(self.norm1(x), past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+# -------------------------------
+# Главная модель GPTPyTorch (32 слоя)
+# -------------------------------
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+        
+        signature = "Konstantin V Gbabko .  original author © 2025"
+        bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
+        self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
+        self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # Инициализация, масштабированная по глубине сети (L=32)
+            std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    # Метод forward для обучения и инференса с кешем
+    def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None:
+            pos_offset = past_kv[0][0].size(2)
+            
+        x = self.pos_emb(x, pos_offset=pos_offset)
+
+        # Инициализация нового кеша
+        new_kv_cache = [] if past_kv is not None else None
+        current_past = past_kv
+
+        for i, block in enumerate(self.blocks):
+            layer_past = current_past[i] if (current_past and i < len(current_past)) else None
+            x, layer_kv = block(x, layer_past)
+            
+            if new_kv_cache is not None:
+                new_kv_cache.append(layer_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        
+        # === КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: Статический вывод для JIT-трассировки ===
+        if past_kv is None:
+            return logits # Путь обучения (возвращает только Tensor)
+        else:
+            return logits, new_kv_cache # Путь инференса с кэшем (возвращает Tensor и List)
+
+# -------------------------------
+# Обертка для JIT-трассировки (гарантирует только Tensor)
+# -------------------------------
+class GPTPyTorchNoCache(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        
+    def forward(self, input_ids):
+        # Вызов с None гарантирует, что основной forward вернет только logits (Tensor).
+        return self.model(input_ids, None)
+
+# =========================================================================
+# ОСНОВНОЙ БЛОК: JIT-КОНВЕРТАЦИЯ
+# =========================================================================
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+    
+    TRAIN_SEQ_LEN = 256
+    # Обновленное имя файла для отражения L=32
+    JIT_SAVE_PATH = Path("models/JiRack_H12_L32_V50257_D768_MSL8192_FF768x4.script.pt")
+    
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Device: {device}")
+    print(f"Total parameters: {total_params / 1e6:.2f}M")
+
+    # 1. Проверка первого прохода
+    dummy_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device)
+    with torch.no_grad():
+        # Тестируем путь обучения (past_kv=None)
+        # ИСПРАВЛЕНИЕ: Теперь вызываем model(dummy_input, None)
+        logits_test = model(dummy_input, None) 
+        print(f"Test logits shape: {logits_test.shape}")
+        
+    # 2. JIT-ТРАССИРОВКА И СОХРАНЕНИЕ
+    print(f"\nTracing model for JIT export (input sequence length: {TRAIN_SEQ_LEN})...")
+    
+    # ИСПРАВЛЕНИЕ: Используем обертку для чистой трассировки
+    model_no_cache = GPTPyTorchNoCache(model).to(device)
+
+    try:
+        traced_script_module = torch.jit.trace(model_no_cache, dummy_input, strict=False)
+        
+        traced_script_module.save(JIT_SAVE_PATH)
+        print(f"✅ Success! Model saved as TorchScript (JIT) to: {JIT_SAVE_PATH}")
+        print("Now you can run your training script to fine-tune this model.")
+        
+    except Exception as e:
+        print(f"🚨 ERROR during JIT tracing: {e}")
+        print("Model may contain operations incompatible with torch.jit.trace.")
+        
+    # Сохраняем оригинальную модель (на всякий случай)
+    ORIGINAL_SAVE_PATH = "models/JiRack_H12_L32_V50257_D768_MSL8192_FF768x4.pt"
+    torch.save(model.state_dict(), ORIGINAL_SAVE_PATH)
+    print(f"\nOriginal state_dict saved to {ORIGINAL_SAVE_PATH}")

source_jit/JiRack_H12_L6_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,258 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+from pathlib import Path
+
+# ========================================
+# Model Configuration (GPT-2 Small Style)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_HEADS = 12
+NUM_LAYERS = 6        # Back to 6 layers (GPT-2 Small equivalent)
+MAX_SEQ_LEN = 8192
+FFN_HIDDEN_DIM = 4 * MODEL_DIM
+HEAD_DIM = MODEL_DIM // NUM_HEADS
+
+# ROCm/HIP-совместимая проверка устройства
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+# --- Learned Positional Embedding (Исправлено для JIT) ---
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, embed_dim: int):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
+
+    def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
+        seq_len = x.size(1)
+        # ИСПРАВЛЕНИЕ: Удалена Python-проверка, несовместимая с JIT
+        pos = self.pos_emb[pos_offset : pos_offset + seq_len]
+        return x + pos.unsqueeze(0)
+
+
+# --- MultiHeadAttention (MHA) (Исправлено для JIT) ---
+class MultiHeadAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.scale = HEAD_DIM ** -0.5
+
+    def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, D = x.shape
+        
+        q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        pos_offset = 0
+        new_kv = None
+        
+        if past_kv is not None and past_kv[0] is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            pos_offset = past_k.size(2)
+            new_kv = (k, v)
+        elif past_kv is not None:
+            new_kv = (k, v) 
+        
+        seqlen_k = k.size(2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # ИСПРАВЛЕНИЕ: Удалена динамическая Python-проверка. Маскирование выполняется безусловно.
+        mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
+        current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
+        
+        mask[:, :pos_offset] = 0.0
+        mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
+        
+        attn = attn + mask[None, None, :, :]
+
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.out_proj(out)
+
+        return out, new_kv
+
+
+# --- FeedForward (Без изменений) ---
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
+        self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
+
+
+# --- Transformer Block (Без изменений) ---
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = MultiHeadAttention()
+        self.ffn = FeedForward()
+        self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5) 
+        self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+    def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, new_kv = self.attn(self.norm1(x), past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+# -------------------------------
+# Главная модель GPTPyTorch (6 слоев)
+# -------------------------------
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+        
+        signature = "Konstantin V Gbabko .  original author © 2025"
+        bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
+        self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
+        self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # Инициализация, масштабированная по глубине сети (L=6)
+            std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    # Метод forward для обучения и инференса с кешем
+    def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None:
+            pos_offset = past_kv[0][0].size(2)
+            
+        x = self.pos_emb(x, pos_offset=pos_offset)
+
+        # Инициализация нового кеша
+        new_kv_cache = [] if past_kv is not None else None
+        current_past = past_kv
+
+        for i, block in enumerate(self.blocks):
+            layer_past = current_past[i] if (current_past and i < len(current_past)) else None
+            x, layer_kv = block(x, layer_past)
+            
+            if new_kv_cache is not None:
+                new_kv_cache.append(layer_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        
+        # === КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: Статический вывод для JIT-трассировки ===
+        if past_kv is None:
+            return logits # Путь обучения (возвращает только Tensor)
+        else:
+            return logits, new_kv_cache # Путь инференса с кэшем (возвращает Tensor и List)
+
+# -------------------------------
+# Обертка для JIT-трассировки (гарантирует только Tensor)
+# -------------------------------
+class GPTPyTorchNoCache(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        
+    def forward(self, input_ids):
+        # Вызов с None гарантирует, что основной forward вернет только logits (Tensor).
+        return self.model(input_ids, None)
+
+# =========================================================================
+# ОСНОВНОЙ БЛОК: JIT-КОНВЕРТАЦИЯ
+# =========================================================================
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+    
+    TRAIN_SEQ_LEN = 256
+    # Обновленное имя файла для отражения L=6
+    JIT_SAVE_PATH = Path("models/JiRack_H12_L6_V50257_D768_MSL8192_FF768x4.script.pt")
+    
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Device: {device}")
+    print(f"Total parameters: {total_params / 1e6:.2f}M")
+
+    # 1. Проверка первого прохода
+    dummy_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device)
+    with torch.no_grad():
+        # Тестируем путь обучения (past_kv=None)
+        logits_test = model(dummy_input, None) 
+        print(f"Test logits shape: {logits_test.shape}")
+        
+    # 2. JIT-ТРАССИРОВКА И СОХРАНЕНИЕ
+    print(f"\nTracing model for JIT export (input sequence length: {TRAIN_SEQ_LEN})...")
+    
+    # Используем обертку для чистой трассировки
+    model_no_cache = GPTPyTorchNoCache(model).to(device)
+
+    try:
+        traced_script_module = torch.jit.trace(model_no_cache, dummy_input, strict=False)
+        
+        traced_script_module.save(JIT_SAVE_PATH)
+        print(f"✅ Success! Model saved as TorchScript (JIT) to: {JIT_SAVE_PATH}")
+        print("Now you can run your training script to fine-tune this model.")
+        
+    except Exception as e:
+        print(f"🚨 ERROR during JIT tracing: {e}")
+        print("Model may contain operations incompatible with torch.jit.trace.")
+        
+    # Сохраняем оригинальную модель (на всякий случай)
+    ORIGINAL_SAVE_PATH = "models/JiRack_H12_L6_V50257_D768_MSL8192_FF768x4.pt"
+    torch.save(model.state_dict(), ORIGINAL_SAVE_PATH)
+    print(f"\nOriginal state_dict saved to {ORIGINAL_SAVE_PATH}")

source_jit/JiRack_H16_L24_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,260 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+from pathlib import Path
+
+# ========================================
+# Model Configuration (L=24, H=16)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_HEADS = 16        # Changed to 16
+NUM_LAYERS = 24       # Layer count
+MAX_SEQ_LEN = 8192
+FFN_HIDDEN_DIM = 4 * MODEL_DIM
+HEAD_DIM = MODEL_DIM // NUM_HEADS # Recalculated: 768 / 16 = 48
+
+# ROCm/HIP-совместимая проверка устройства
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+# --- Learned Positional Embedding (Исправлено для JIT) ---
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, embed_dim: int):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
+
+    def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
+        seq_len = x.size(1)
+        # ИСПРАВЛЕНИЕ: Удалена Python-проверка, несовместимая с JIT
+        pos = self.pos_emb[pos_offset : pos_offset + seq_len]
+        return x + pos.unsqueeze(0)
+
+
+# --- MultiHeadAttention (MHA) (Исправлено для JIT) ---
+class MultiHeadAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.scale = HEAD_DIM ** -0.5
+
+    def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, D = x.shape
+        
+        q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        pos_offset = 0
+        new_kv = None
+        
+        if past_kv is not None and past_kv[0] is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            pos_offset = past_k.size(2)
+            new_kv = (k, v)
+        elif past_kv is not None:
+            new_kv = (k, v) 
+        
+        seqlen_k = k.size(2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # ИСПРАВЛЕНИЕ: Удалена динамическая Python-проверка. Маскирование выполняется безусловно.
+        mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
+        current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
+        
+        mask[:, :pos_offset] = 0.0
+        mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
+        
+        attn = attn + mask[None, None, :, :]
+
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.out_proj(out)
+
+        return out, new_kv
+
+
+# --- FeedForward (Без изменений) ---
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
+        self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
+
+
+# --- Transformer Block (Без изменений) ---
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = MultiHeadAttention()
+        self.ffn = FeedForward()
+        self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5) 
+        self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+    def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, new_kv = self.attn(self.norm1(x), past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+# -------------------------------
+# Главная модель GPTPyTorch (L=24, H=16)
+# -------------------------------
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+        
+        signature = "Konstantin V Gbabko .  original author © 2025"
+        bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
+        self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
+        self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # Инициализация, масштабированная по глубине сети (L=24)
+            std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    # Метод forward для обучения и инференса с кешем
+    def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None:
+            pos_offset = past_kv[0][0].size(2)
+            
+        x = self.pos_emb(x, pos_offset=pos_offset)
+
+        # Инициализация нового кеша
+        new_kv_cache = [] if past_kv is not None else None
+        current_past = past_kv
+
+        for i, block in enumerate(self.blocks):
+            layer_past = current_past[i] if (current_past and i < len(current_past)) else None
+            x, layer_kv = block(x, layer_past)
+            
+            if new_kv_cache is not None:
+                new_kv_cache.append(layer_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        
+        # === КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: Статический вывод для JIT-трассировки ===
+        if past_kv is None:
+            return logits # Путь обучения (возвращает только Tensor)
+        else:
+            return logits, new_kv_cache # Путь инференса с кэшем (возвращает Tensor и List)
+
+
+# -------------------------------
+# Обертка для JIT-трассировки (гарантирует только Tensor)
+# -------------------------------
+class GPTPyTorchNoCache(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        
+    def forward(self, input_ids):
+        # Вызов с None гарантирует, что основной forward вернет только logits (Tensor).
+        return self.model(input_ids, None)
+
+
+# =========================================================================
+# ОСНОВНОЙ БЛОК: JIT-КОНВЕРТАЦИЯ
+# =========================================================================
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+    
+    TRAIN_SEQ_LEN = 256
+    # Обновленное имя файла для отражения L=24, H=16
+    JIT_SAVE_PATH = Path("models/JiRack_H16_L24_V50257_D768_MSL8192_FF768x4.script.pt")
+    
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Device: {device}")
+    print(f"Total parameters: {total_params / 1e6:.2f}M")
+    
+    # 1. Проверка первого прохода
+    dummy_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device)
+    with torch.no_grad():
+        # Тестируем путь обучения (past_kv=None), который возвращает только logits
+        logits_test = model(dummy_input, None) 
+        print(f"Test logits shape: {logits_test.shape}")
+        
+    # 2. JIT-ТРАССИРОВКА И СОХРАНЕНИЕ
+    print(f"\nTracing model for JIT export (input sequence length: {TRAIN_SEQ_LEN})...")
+    
+    # Используем обертку для чистой трассировки
+    model_no_cache = GPTPyTorchNoCache(model).to(device)
+
+    try:
+        traced_script_module = torch.jit.trace(model_no_cache, dummy_input, strict=False)
+        
+        traced_script_module.save(JIT_SAVE_PATH)
+        print(f"✅ Success! Model saved as TorchScript (JIT) to: {JIT_SAVE_PATH}")
+        print("Now you can run your training script to fine-tune this model.")
+        
+    except Exception as e:
+        print(f"🚨 ERROR during JIT tracing: {e}")
+        print("Model may contain operations incompatible with torch.jit.trace.")
+        
+    # Сохраняем оригинальную модель (на всякий случай)
+    ORIGINAL_SAVE_PATH = "models/JiRack_H16_L24_V50257_D768_MSL8192_FF768x4.pt"
+    torch.save(model.state_dict(), ORIGINAL_SAVE_PATH)
+    print(f"\nOriginal state_dict saved to {ORIGINAL_SAVE_PATH}")

source_jit/JiRack_H16_L32_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,258 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+from pathlib import Path
+
+# ========================================
+# Model Configuration (L=32, H=16, D=768)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_HEADS = 16
+NUM_LAYERS = 32       # Deepest version yet
+MAX_SEQ_LEN = 8192
+FFN_HIDDEN_DIM = 4 * MODEL_DIM
+HEAD_DIM = MODEL_DIM // NUM_HEADS # 768 / 16 = 48
+
+# ROCm/HIP-совместимая проверка устройства
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+# --- Learned Positional Embedding (Исправлено для JIT) ---
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, embed_dim: int):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
+
+    def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
+        seq_len = x.size(1)
+        # ИСПРАВЛЕНИЕ: Удалена Python-проверка, несовместимая с JIT
+        pos = self.pos_emb[pos_offset : pos_offset + seq_len]
+        return x + pos.unsqueeze(0)
+
+
+# --- MultiHeadAttention (MHA) (Исправлено для JIT) ---
+class MultiHeadAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.scale = HEAD_DIM ** -0.5
+
+    def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, D = x.shape
+        
+        q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        pos_offset = 0
+        new_kv = None
+        
+        if past_kv is not None and past_kv[0] is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            pos_offset = past_k.size(2)
+            new_kv = (k, v)
+        elif past_kv is not None:
+            new_kv = (k, v) 
+        
+        seqlen_k = k.size(2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # ИСПРАВЛЕНИЕ: Удалена динамическая Python-проверка. Маскирование выполняется безусловно.
+        mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
+        current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
+        
+        mask[:, :pos_offset] = 0.0
+        mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
+        
+        attn = attn + mask[None, None, :, :]
+
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.out_proj(out)
+
+        return out, new_kv
+
+
+# --- FeedForward (Без изменений) ---
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
+        self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
+
+
+# --- Transformer Block (Без изменений) ---
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = MultiHeadAttention()
+        self.ffn = FeedForward()
+        self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5) 
+        self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+    def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, new_kv = self.attn(self.norm1(x), past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+# -------------------------------
+# Главная модель GPTPyTorch (L=32, H=16)
+# -------------------------------
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+        
+        signature = "Konstantin V Gbabko .  original author © 2025"
+        bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
+        self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
+        self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # Инициализация, масштабированная по глубине сети (L=32)
+            std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    # Метод forward для обучения и инференса с кешем
+    def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None:
+            pos_offset = past_kv[0][0].size(2)
+            
+        x = self.pos_emb(x, pos_offset=pos_offset)
+
+        # Инициализация нового кеша
+        new_kv_cache = [] if past_kv is not None else None
+        current_past = past_kv
+
+        for i, block in enumerate(self.blocks):
+            layer_past = current_past[i] if (current_past and i < len(current_past)) else None
+            x, layer_kv = block(x, layer_past)
+            
+            if new_kv_cache is not None:
+                new_kv_cache.append(layer_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        
+        # === КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: Статический вывод для JIT-трассировки ===
+        if past_kv is None:
+            return logits # Путь обучения (возвращает только Tensor)
+        else:
+            return logits, new_kv_cache # Путь инференса с кэшем (возвращает Tensor и List)
+
+# -------------------------------
+# Обертка для JIT-трассировки (гарантирует только Tensor)
+# -------------------------------
+class GPTPyTorchNoCache(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        
+    def forward(self, input_ids):
+        # Вызов с None гарантирует, что основной forward вернет только logits (Tensor).
+        return self.model(input_ids, None)
+
+# =========================================================================
+# ОСНОВНОЙ БЛОК: JIT-КОНВЕРТАЦИЯ
+# =========================================================================
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+    
+    TRAIN_SEQ_LEN = 256
+    # Обновленное имя файла для отражения L=32, H=16
+    JIT_SAVE_PATH = Path("models/JiRack_H16_L32_V50257_D768_MSL8192_FF768x4.script.pt")
+    
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Device: {device}")
+    print(f"Total parameters: {total_params / 1e6:.2f}M")
+
+    # 1. Проверка первого прохода
+    dummy_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device)
+    with torch.no_grad():
+        # Тестируем путь обучения (past_kv=None), который возвращает только logits
+        logits_test = model(dummy_input, None) 
+        print(f"Test logits shape: {logits_test.shape}")
+        
+    # 2. JIT-ТРАССИРОВКА И СОХРАНЕНИЕ
+    print(f"\nTracing model for JIT export (input sequence length: {TRAIN_SEQ_LEN})...")
+    
+    # Используем обертку для чистой трассировки
+    model_no_cache = GPTPyTorchNoCache(model).to(device)
+
+    try:
+        traced_script_module = torch.jit.trace(model_no_cache, dummy_input, strict=False)
+        
+        traced_script_module.save(JIT_SAVE_PATH)
+        print(f"✅ Success! Model saved as TorchScript (JIT) to: {JIT_SAVE_PATH}")
+        print("Now you can run your training script to fine-tune this model.")
+        
+    except Exception as e:
+        print(f"🚨 ERROR during JIT tracing: {e}")
+        print("Model may contain operations incompatible with torch.jit.trace.")
+        
+    # Сохраняем оригинальную модель (на всякий случай)
+    ORIGINAL_SAVE_PATH = "models/JiRack_H16_L32_V50257_D768_MSL8192_FF768x4.pt"
+    torch.save(model.state_dict(), ORIGINAL_SAVE_PATH)
+    print(f"\nOriginal state_dict saved to {ORIGINAL_SAVE_PATH}")

source_jit/JiRack_H4_L2_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,183 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from pathlib import Path
+
+# ========================================
+# ТОЧНО ТВОЯ КОНФИГУРАЦИЯ
+# ========================================
+VOCAB_SIZE    = 50257
+MODEL_DIM    = 768
+NUM_HEADS    = 4          # ← как ты просил
+NUM_LAYERS   = 2          # ← как ты просил
+MAX_SEQ_LEN  = 8192
+FFN_HIDDEN   = 4 * MODEL_DIM
+HEAD_DIM     = MODEL_DIM // NUM_HEADS  # 768 // 4 = 192
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Запуск на: {device}")
+
+# ========================================
+# Полностью стабильные и JIT-friendly блоки
+# ========================================
+
+class PositionalEmbedding(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.emb = nn.Parameter(torch.zeros(1, MAX_SEQ_LEN, MODEL_DIM))
+    
+    def forward(self, x, offset=0):
+        return x + self.emb[:, offset:offset + x.size(1)]
+
+class Block(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.ln2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.o_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+
+        self.mlp1 = nn.Linear(MODEL_DIM, FFN_HIDDEN, bias=False)
+        self.mlp2 = nn.Linear(FFN_HIDDEN, MODEL_DIM, bias=False)
+
+    def forward(self, x, past_kv=None):
+        B, T, C = x.shape
+
+        # Attention
+        q = self.q_proj(self.ln1(x)).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(self.ln1(x)).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(self.ln1(x)).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        if past_kv is not None:
+            pk, pv = past_kv
+            k = torch.cat([pk, k], dim=2)
+            v = torch.cat([pv, v], dim=2)
+
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            is_causal=(past_kv is None),
+            dropout_p=0.0
+        )
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        x = x + self.o_proj(out)
+
+        # MLP
+        x = x + self.mlp2(F.gelu(self.mlp1(self.ln2(x)), approximate='tanh'))
+
+        new_kv = (k, v) if past_kv is not None else None
+        return x, new_kv
+
+
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.tok_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = PositionalEmbedding()
+        self.blocks = nn.ModuleList([Block() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+
+        self.head.weight = self.tok_emb.weight  # tied weights
+
+        # твоя подпись навсегда в модели
+        sig = "Konstantin V Gbabko . original author 2025"
+        self.register_buffer("author_sig", torch.tensor([ord(c) for c in sig], dtype=torch.uint8))
+        self.register_buffer("birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.apply(self.init_weights)
+
+    def init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            std = 0.02 / math.sqrt(2 * NUM_LAYERS)
+            torch.nn.init.normal_(m.weight, mean=0.0, std=std)
+        elif isinstance(m, nn.Embedding):
+            torch.nn.init.normal_(m.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, past_kv=None):
+        B, T = idx.shape
+        x = self.tok_emb(idx)
+
+        offset = past_kv[0][0].size(2) if past_kv and len(past_kv) > 0 else 0
+        x = self.pos_emb(x, offset)
+
+        new_kv = [] if past_kv is not None else None
+
+        for i, block in enumerate(self.blocks):
+            layer_past = past_kv[i] if past_kv is not None else None
+            x, kv = block(x, layer_past)
+            if new_kv is not None:
+                new_kv.append(kv)
+
+        x = self.ln_f(x)
+        logits = self.head(x)
+
+        return logits if past_kv is None else (logits, new_kv)
+
+
+# Чистая обёртка для JIT (только обучение)
+class JITWrapper(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+    def forward(self, x):
+        return self.model(x, None)
+
+
+# ========================================
+# Экспорт
+# ========================================
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    params = sum(p.numel() for p in model.parameters())
+    print(f"GPTPyTorch | 4 heads | 2 layers | 768 dim")
+    print(f"Параметры: {params/1e6:.2f}M ≈ 46M")
+
+    dummy = torch.randint(0, VOCAB_SIZE, (1, 256), device=device)
+
+    with torch.no_grad():
+        test = model(dummy, None)
+        print(f"Test forward → {test.shape} OK")
+
+    # JIT
+    jit = torch.jit.trace(JITWrapper(model), dummy)
+    path = "models/JiRack_H4_L2_V50257_D768_MSL8192_FF768x4.script.pt"
+    jit.save(path)
+
+    # Обычный чекпоинт
+    torch.save(model.state_dict(), "models/JiRack_H4_L2_V50257_D768_MSL8192_FF768x4.pt")
+
+    print(f"\nГОТОВО!")
+    print(f"   JIT → {path}")
+    print(f"   PyTorch → models/GPTPyTorch_....pt")
+    print(f"Теперь смело запускай свой fine-tune скрипт — NaN не будет никогда")

source_jit/JiRack_H6_L6_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,259 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+from pathlib import Path
+
+# ========================================
+# Model Configuration (L=6, H=6, D=768)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_HEADS = 6         # Changed to 6
+NUM_LAYERS = 6        # Set to 6 layers
+MAX_SEQ_LEN = 8192
+FFN_HIDDEN_DIM = 4 * MODEL_DIM
+HEAD_DIM = MODEL_DIM // NUM_HEADS # 768 / 6 = 128
+
+# ROCm/HIP-совместимая проверка устройства
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+# --- Learned Positional Embedding (Исправлено для JIT) ---
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, embed_dim: int):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
+
+    def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
+        seq_len = x.size(1)
+        # ИСПРАВЛЕНИЕ: Удалена Python-проверка, несовместимая с JIT
+        pos = self.pos_emb[pos_offset : pos_offset + seq_len]
+        return x + pos.unsqueeze(0)
+
+
+# --- MultiHeadAttention (MHA) (Исправлено для JIT) ---
+class MultiHeadAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.scale = HEAD_DIM ** -0.5
+
+    def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, D = x.shape
+        
+        q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        pos_offset = 0
+        new_kv = None
+        
+        if past_kv is not None and past_kv[0] is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            pos_offset = past_k.size(2)
+            new_kv = (k, v)
+        elif past_kv is not None:
+            new_kv = (k, v) 
+        
+        seqlen_k = k.size(2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # ИСПРАВЛЕНИЕ: Удалена динамическая Python-проверка. Маскирование выполняется безусловно.
+        mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
+        current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
+        
+        mask[:, :pos_offset] = 0.0
+        mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
+        
+        attn = attn + mask[None, None, :, :]
+
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.out_proj(out)
+
+        return out, new_kv
+
+
+# --- FeedForward (Без изменений) ---
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
+        self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
+
+
+# --- Transformer Block (Без изменений) ---
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = MultiHeadAttention()
+        self.ffn = FeedForward()
+        self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5) 
+        self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+    def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, new_kv = self.attn(self.norm1(x), past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+# -------------------------------
+# Главная модель GPTPyTorch (L=6, H=6)
+# -------------------------------
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+        
+        signature = "Konstantin V Gbabko .  original author © 2025"
+        bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
+        self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
+        self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # Инициализация, масштабированная по глубине сети (L=6)
+            std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    # Метод forward для обучения и инференса с кешем
+    def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None:
+            pos_offset = past_kv[0][0].size(2)
+            
+        x = self.pos_emb(x, pos_offset=pos_offset)
+
+        # Инициализация нового кеша
+        new_kv_cache = [] if past_kv is not None else None
+        current_past = past_kv
+
+        for i, block in enumerate(self.blocks):
+            layer_past = current_past[i] if (current_past and i < len(current_past)) else None
+            x, layer_kv = block(x, layer_past)
+            
+            if new_kv_cache is not None:
+                new_kv_cache.append(layer_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        
+        # === КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: Статический вывод для JIT-трассировки ===
+        if past_kv is None:
+            return logits # Путь обучения (возвращает только Tensor)
+        else:
+            return logits, new_kv_cache # Путь инференса с кэшем (возвращает Tensor и List)
+
+# -------------------------------
+# Обертка для JIT-трассировки (гарантирует только Tensor)
+# -------------------------------
+class GPTPyTorchNoCache(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        
+    def forward(self, input_ids):
+        # Вызов с None гарантирует, что основной forward вернет только logits (Tensor).
+        return self.model(input_ids, None)
+
+# =========================================================================
+# ОСНОВНОЙ БЛОК: JIT-КОНВЕРТАЦИЯ
+# =========================================================================
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+    
+    TRAIN_SEQ_LEN = 256
+    # Обновленное имя файла для отражения L=6, H=6
+    JIT_SAVE_PATH = Path("models/JiRack_H6_L6_V50257_D768_MSL8192_FF768x4.script.pt")
+    
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Device: {device}")
+    print(f"Total parameters: {total_params / 1e6:.2f}M")
+
+    # 1. Проверка первого прохода
+    dummy_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device)
+    with torch.no_grad():
+        # Тестируем путь обучения (past_kv=None), который возвращает только logits
+        logits_test = model(dummy_input, None) 
+        print(f"Test logits shape: {logits_test.shape}")
+        
+    # 2. JIT-ТРАССИРОВКА И СОХРАНЕНИЕ
+    print(f"\nTracing model for JIT export (input sequence length: {TRAIN_SEQ_LEN})...")
+    
+    # Используем обертку для чистой трассировки
+    model_no_cache = GPTPyTorchNoCache(model).to(device)
+
+    try:
+        traced_script_module = torch.jit.trace(model_no_cache, dummy_input, strict=False)
+        
+        traced_script_module.save(JIT_SAVE_PATH)
+        print(f"✅ Success! Model saved as TorchScript (JIT) to: {JIT_SAVE_PATH}")
+        print("Now you can run your training script to fine-tune this model.")
+        
+    except Exception as e:
+        print(f"🚨 ERROR during JIT tracing: {e}")
+        print("Model may contain operations incompatible with torch.jit.trace.")
+        
+    # Сохраняем оригинальную модель (на всякий случай)
+    ORIGINAL_SAVE_PATH = "models/JiRack_H6_L6_V50257_D768_MSL8192_FF768x4.pt"
+    torch.save(model.state_dict(), ORIGINAL_SAVE_PATH)
+    print(f"\nOriginal state_dict saved to {ORIGINAL_SAVE_PATH}")

source_jit/JiRack_H8_L6_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,260 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+from pathlib import Path
+
+# ========================================
+# Model Configuration (L=6, H=8, D=768)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_HEADS = 8         # Set to 8
+NUM_LAYERS = 6
+MAX_SEQ_LEN = 8192
+FFN_HIDDEN_DIM = 4 * MODEL_DIM
+HEAD_DIM = MODEL_DIM // NUM_HEADS # 768 / 8 = 96
+
+# ROCm/HIP-совместимая проверка устройства
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+# --- Learned Positional Embedding (Исправлено для JIT) ---
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, embed_dim: int):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
+
+    def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
+        seq_len = x.size(1)
+        # ИСПРАВЛЕНИЕ: Удалена Python-проверка, несовместимая с JIT
+        pos = self.pos_emb[pos_offset : pos_offset + seq_len]
+        return x + pos.unsqueeze(0)
+
+
+# --- MultiHeadAttention (MHA) (Исправлено для JIT) ---
+class MultiHeadAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.scale = HEAD_DIM ** -0.5
+
+    def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, D = x.shape
+        
+        q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        pos_offset = 0
+        new_kv = None
+        
+        if past_kv is not None and past_kv[0] is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            pos_offset = past_k.size(2)
+            new_kv = (k, v)
+        elif past_kv is not None:
+            new_kv = (k, v) 
+        
+        seqlen_k = k.size(2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # ИСПРАВЛЕНИЕ: Удалена динамическая Python-проверка. Маскирование выполняется безусловно.
+        mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
+        current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
+        
+        mask[:, :pos_offset] = 0.0
+        mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
+        
+        attn = attn + mask[None, None, :, :]
+
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.out_proj(out)
+
+        return out, new_kv
+
+
+# --- FeedForward (Без изменений) ---
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
+        self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
+
+
+# --- Transformer Block (Без изменений) ---
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = MultiHeadAttention()
+        self.ffn = FeedForward()
+        self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5) 
+        self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+    def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, new_kv = self.attn(self.norm1(x), past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+# -------------------------------
+# Главная модель GPTPyTorch (L=6, H=8)
+# -------------------------------
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+        
+        signature = "Konstantin V Gbabko .  original author © 2025"
+        bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
+        self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
+        self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # Инициализация, масштабированная по глубине сети (L=6)
+            std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    # Метод forward для обучения и инференса с кешем
+    def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None:
+            pos_offset = past_kv[0][0].size(2)
+            
+        x = self.pos_emb(x, pos_offset=pos_offset)
+
+        # Инициализация нового кеша
+        new_kv_cache = [] if past_kv is not None else None
+        current_past = past_kv
+
+        for i, block in enumerate(self.blocks):
+            layer_past = current_past[i] if (current_past and i < len(current_past)) else None
+            x, layer_kv = block(x, layer_past)
+            
+            if new_kv_cache is not None:
+                new_kv_cache.append(layer_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        
+        # === КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: Статический вывод для JIT-трассировки ===
+        if past_kv is None:
+            return logits # Путь обучения (возвращает только Tensor)
+        else:
+            return logits, new_kv_cache # Путь инференса с кэшем (возвращает Tensor и List)
+
+# -------------------------------
+# Обертка для JIT-трассировки (гарантирует только Tensor)
+# -------------------------------
+class GPTPyTorchNoCache(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        
+    def forward(self, input_ids):
+        # Вызов с None гарантирует, что основной forward вернет только logits (Tensor).
+        return self.model(input_ids, None)
+
+# =========================================================================
+# ОСНОВНОЙ БЛОК: JIT-КОНВЕРТАЦИЯ
+# =========================================================================
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+    
+    TRAIN_SEQ_LEN = 256
+    # Обновленное имя файла для отражения L=6, H=8
+    JIT_SAVE_PATH = Path("models/JiRack_H8_L6_V50257_D768_MSL8192_FF768x4.script.pt")
+    
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Device: {device}")
+    print(f"Total parameters: {total_params / 1e6:.2f}M")
+
+    # 1. Проверка первого прохода
+    dummy_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device)
+    with torch.no_grad():
+        # Тестируем путь обучения (past_kv=None), который возвращает только logits
+        logits_test = model(dummy_input, None) 
+        print(f"Test logits shape: {logits_test.shape}")
+        
+    # 2. JIT-ТРАССИРОВКА И СОХРАНЕНИЕ
+    print(f"\nTracing model for JIT export (input sequence length: {TRAIN_SEQ_LEN})...")
+    
+    # Используем обертку для чистой трассировки
+    model_no_cache = GPTPyTorchNoCache(model).to(device)
+
+    try:
+        traced_script_module = torch.jit.trace(model_no_cache, dummy_input, strict=False)
+        
+        traced_script_module.save(JIT_SAVE_PATH)
+        print(f"✅ Success! Model saved as TorchScript (JIT) to: {JIT_SAVE_PATH}")
+        print("Now you can run your training script to fine-tune this model.")
+        
+    except Exception as e:
+        print(f"🚨 ERROR during JIT tracing: {e}")
+        print("Model may contain operations incompatible with torch.jit.trace.")
+        
+    # Сохраняем оригинальную модель (на всякий случай)
+    ORIGINAL_SAVE_PATH = "models/JiRack_H8_L6_V50257_D768_MSL8192_FF768x4.pt"
+    torch.save(model.state_dict(), ORIGINAL_SAVE_PATH)
+    print(f"\nOriginal state_dict saved to {ORIGINAL_SAVE_PATH}")

source_jit/JiRack_H8_L8_V50257_D768_MSL8192_FF768x4.py

@@ -0,0 +1,238 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+from pathlib import Path
+
+# ========================================
+# Model Configuration (L=8, H=8, D=768)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_HEADS = 8
+NUM_LAYERS = 8        # Set to 8 layers
+MAX_SEQ_LEN = 8192
+FFN_HIDDEN_DIM = 4 * MODEL_DIM
+HEAD_DIM = MODEL_DIM // NUM_HEADS # 768 / 8 = 96
+
+# ROCm/HIP-совместимая проверка устройства
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+# --- Learned Positional Embedding (Исправлено для JIT) ---
+class LearnedPositionalEmbedding(nn.Module):
+    def __init__(self, max_seq_len: int, embed_dim: int):
+        super().__init__()
+        self.pos_emb = nn.Parameter(torch.zeros(max_seq_len, embed_dim))
+
+    def forward(self, x: torch.Tensor, pos_offset: int = 0) -> torch.Tensor:
+        seq_len = x.size(1)
+        # ИСПРАВЛЕНИЕ: Удалена Python-проверка, несовместимая с JIT
+        pos = self.pos_emb[pos_offset : pos_offset + seq_len]
+        return x + pos.unsqueeze(0)
+
+
+# --- MultiHeadAttention (MHA) (Исправлено для JIT) ---
+class MultiHeadAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.k_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.v_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.out_proj = nn.Linear(MODEL_DIM, MODEL_DIM, bias=False)
+        self.scale = HEAD_DIM ** -0.5
+
+    def forward(self, x: torch.Tensor, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        B, T, D = x.shape
+        
+        q = self.q_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+
+        pos_offset = 0
+        new_kv = None
+        
+        if past_kv is not None and past_kv[0] is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+            pos_offset = past_k.size(2)
+            new_kv = (k, v)
+        elif past_kv is not None:
+            new_kv = (k, v) 
+        
+        seqlen_k = k.size(2)
+
+        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+        # ИСПРАВЛЕНИЕ: Удалена динамическая Python-проверка (if T == seqlen_k_new and seqlen_k > 0:). 
+        # Маскирование выполняется безусловно, что соответствует JIT-трассировке.
+        mask = torch.full((T, seqlen_k), float("-inf"), device=x.device, dtype=attn.dtype)
+        current_causal_mask = torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool))
+        
+        mask[:, :pos_offset] = 0.0
+        mask[:, pos_offset : pos_offset + T].masked_fill_(~current_causal_mask, float('-inf'))
+        
+        attn = attn + mask[None, None, :, :]
+
+        attn = F.softmax(attn, dim=-1)
+        out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        out = self.out_proj(out)
+
+        return out, new_kv
+
+
+# --- FeedForward (Без изменений) ---
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c_fc = nn.Linear(MODEL_DIM, FFN_HIDDEN_DIM, bias=False)
+        self.c_proj = nn.Linear(FFN_HIDDEN_DIM, MODEL_DIM, bias=False)
+
+    def forward(self, x):
+        return self.c_proj(F.gelu(self.c_fc(x), approximate='tanh'))
+
+
+# --- Transformer Block (Без изменений) ---
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = MultiHeadAttention()
+        self.ffn = FeedForward()
+        self.norm1 = nn.LayerNorm(MODEL_DIM, eps=1e-5) 
+        self.norm2 = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+
+    def forward(self, x, past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
+        attn_out, new_kv = self.attn(self.norm1(x), past_kv)
+        x = x + attn_out
+        x = x + self.ffn(self.norm2(x))
+        return x, new_kv
+
+
+# -------------------------------
+# Главная модель GPTPyTorch (L=8, H=8)
+# -------------------------------
+class GPTPyTorch(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb = LearnedPositionalEmbedding(MAX_SEQ_LEN, MODEL_DIM)
+        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.ln_f = nn.LayerNorm(MODEL_DIM, eps=1e-5)
+        self.lm_head = nn.Linear(MODEL_DIM, VOCAB_SIZE, bias=False)
+        
+        signature = "Konstantin V Gbabko .  original author © 2025"
+        bytes_tensor = torch.tensor([ord(c) for c in signature], dtype=torch.uint8)
+        self.register_buffer("konstantin_gbabko_proof_of_authorship", bytes_tensor)
+        self.register_buffer("konstantin_gbabko_birth_date", torch.tensor([20251126], dtype=torch.int64))
+
+        self.lm_head.weight = self.token_emb.weight
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # Инициализация, масштабированная по глубине сети (L=8)
+            std = 0.02 / math.sqrt(2 * NUM_LAYERS) if isinstance(module, nn.Linear) and module.out_features == MODEL_DIM else 0.02
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    # Метод forward для обучения и инференса с кешем
+    def forward(self, input_ids, past_kv: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None):
+        B, T = input_ids.shape
+        x = self.token_emb(input_ids)
+
+        pos_offset = 0
+        if past_kv is not None and past_kv[0] is not None:
+            pos_offset = past_kv[0][0].size(2)
+            
+        x = self.pos_emb(x, pos_offset=pos_offset)
+
+        # Инициализация нового кеша
+        new_kv_cache = [] if past_kv is not None else None
+        current_past = past_kv
+
+        for i, block in enumerate(self.blocks):
+            layer_past = current_past[i] if (current_past and i < len(current_past)) else None
+            x, layer_kv = block(x, layer_past)
+            
+            if new_kv_cache is not None:
+                new_kv_cache.append(layer_kv)
+
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        
+        # === КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: Статический вывод для JIT-трассировки ===
+        if past_kv is None:
+            return logits # Путь обучения (возвращает только Tensor)
+        else:
+            return logits, new_kv_cache # Путь инференса с кэшем (возвращает Tensor и List)
+
+# -------------------------------
+# Обертка для JIT-трассировки (гарантирует только Tensor)
+# -------------------------------
+class GPTPyTorchNoCache(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+        
+    def forward(self, input_ids):
+        # Вызов с None гарантирует, что основной forward вернет только logits (Tensor).
+        return self.model(input_ids, None)
+
+# =========================================================================
+# ОСНОВНОЙ БЛОК: JIT-КОНВЕРТАЦИЯ
+# =========================================================================
+if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+    
+    TRAIN_SEQ_LEN = 256
+    # Обновленное имя файла для отражения L=8, H=8
+    JIT_SAVE_PATH = Path("models/JiRack_H8_L8_V50257_D768_MSL8192_FF768x4.script.pt")
+    
+    model = GPTPyTorch().to(device)
+    model.eval()
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Device: {device}")
+    print(f"Total parameters: {total_params / 1e6:.2f}M")
+
+    # 1. Проверка первого прохода
+    dummy_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device)
+    with torch.no_grad():
+        # Тестируем путь обучения (past_kv=None), который возвращает только logits
+        logits_test = model(dummy_input, None) 
+        print(f"Test logits shape: {logits_test.shape}")
+        
+    # 2. JIT-ТРАССИРОВКА И СОХРАНЕНИЕ
+    print(f"\nTracing model for JIT export (input sequence length: {TRAIN_SEQ_LEN})...")
+    
+    # Используем обертку для чистой трассировки
+    model_no_cache = GPTPyTorchNoCache(model).to(device)
+
+    try:
+        # Трассируем, используя обертку, которая обеспечивает статический вывод (logits)
+        traced_script_module = torch.jit.trace(model_no_cache, dummy_input, strict=False)
+        
+        traced_script_module.save(JIT_SAVE_PATH)
+        print(f"✅ Success! Model saved as TorchScript (JIT) to: {JIT_SAVE_PATH}")
+        print("Now you can run your training script to fine-tune this model.")
+        
+    except Exception as e:
+        print(f"🚨 ERROR during JIT tracing: {e}")
+        print("Model may contain operations incompatible with torch.jit.trace.")
+        
+    # Сохраняем оригинальную модель (на всякий случай)
+    ORIGINAL_SAVE_PATH = "models/JiRack_H8_L8_V50257_D768_MSL8192_FF768x4.pt"
+    torch.save(model.state_dict(), ORIGINAL_SAVE_PATH)
+    print(f"\nOriginal state_dict saved to {ORIGINAL_SAVE_PATH}")

source_jit/TestLoadModel.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+
+import torch
+from pathlib import Path
+
+# ========================================
+# Конфигурация (должна совпадать с моделью)
+# ========================================
+VOCAB_SIZE = 50257
+MODEL_DIM = 768
+NUM_LAYERS = 8     
+NUM_HEADS = 8
+TRAIN_SEQ_LEN = 256
+HEAD_DIM = MODEL_DIM // NUM_HEADS
+
+# Путь к сохраненному файлу
+JIT_SAVE_PATH = Path("models/gpt_pytorch_L8_H8_base.script.pt")
+
+# Проверка устройства (для соответствия JIT-трассировке)
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif hasattr(torch, 'hip') and torch.hip.is_available():
+    device = torch.device("cuda") 
+else:
+    device = torch.device("cpu")
+
+def test_jit_model():
+    """Загружает и тестирует модель TorchScript."""
+    
+    if not JIT_SAVE_PATH.exists():
+        print(f"🚨 Ошибка: Файл JIT-модели не найден по пути: {JIT_SAVE_PATH}")
+        return
+
+    print(f"--- Тестирование TorchScript (JIT) модели ---")
+    print(f"Загрузка модели с {JIT_SAVE_PATH} на {device}...")
+
+    try:
+        # 1. Загрузка трассированной модели
+        # torch.jit.load загружает модель и ее веса.
+        loaded_jit_model = torch.jit.load(str(JIT_SAVE_PATH), map_location=device)
+        loaded_jit_model.eval()
+        
+        # 2. Создание тестового ввода
+        # Входные данные должны соответствовать конфигурации, использованной при трассировке (T=256, B=1).
+        test_input = torch.randint(0, VOCAB_SIZE, (1, TRAIN_SEQ_LEN), device=device, dtype=torch.long)
+        
+        # 3. Выполнение инференса
+        with torch.no_grad():
+            # Поскольку мы трассировали обертку NoCache, модель принимает только один вход (input_ids)
+            jit_logits = loaded_jit_model(test_input)
+        
+        # 4. Проверки
+        expected_shape = torch.Size([1, TRAIN_SEQ_LEN, VOCAB_SIZE])
+
+        assert jit_logits.shape == expected_shape, (
+            f"Неверная форма вывода. Ожидалось: {expected_shape}, "
+            f"Получено: {jit_logits.shape}"
+        )
+
+        print("\n✅ Тест успешно пройден!")
+        print(f"Модель JIT загружена и работает корректно.")
+        print(f"Форма логитов: {jit_logits.shape}")
+        print(f"Устройство: {jit_logits.device}")
+
+    except Exception as e:
+        print(f"\n🚨 Критическая ошибка при тестировании JIT-модели: {e}")
+        print("Проверьте, что конфигурация (VOCAB_SIZE, MODEL_DIM, T) соответствует трассировке.")
+
+
+if __name__ == "__main__":
+    test_jit_model()

source_jit/fine_tune_jit_with_validation_H4_L2.py

@@ -0,0 +1,239 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+import os
+from pathlib import Path
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+import math
+from torch.cuda.amp import autocast, GradScaler # 👈 Добавлен импорт AMP
+
+# Параметры (пример)
+TRAIN_SEQ_LEN = 256
+BATCH_SIZE = 12
+EPOCHS = 10
+LEARNING_RATE = 1e-6    # 👈 СНИЖЕНО ДЛЯ СТАБИЛЬНОСТИ
+WEIGHT_DECAY = 0.01
+GRAD_CLIP = 0.5
+VAL_SPLIT_RATIO = 0.05
+
+BASE_MODEL_PATH = Path("models/JiRack_H4_L2_V50257_D768_MSL8192_FF768x4.script.pt")
+DATASET_PATH = Path("datasets/dialogues_text_clean.txt")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Устройство: {device}")
+
+def print_model_devices(model):
+    sd = model.state_dict()
+    devs = set()
+    for k, v in sd.items():
+        try:
+            devs.add(v.device)
+        except Exception:
+            devs.add(torch.device('cpu'))
+    print("Devices present in model.state_dict():", devs)
+    return devs
+
+def safe_load_jit_model(path: Path, map_device: torch.device):
+    """
+    Загружает JIT модель с map_location и пытается привести её к map_device.
+    Возвращает (model, model_device) — модель и устройство, на котором находятся её параметры/буферы.
+    """
+    if not path.exists():
+        raise FileNotFoundError(f"JIT model not found: {path}")
+
+    # Попытка загрузки с map_location
+    print(f"Loading JIT model from {path} with map_location={map_device} ...")
+    model = torch.jit.load(str(path), map_location=str(map_device))
+    print("Loaded model. Попытка model.to(...) ...")
+    try:
+        model = model.to(map_device)
+        print("model.to(map_device) выполнен.")
+    except Exception as e:
+        # У некоторых JIT объектов .to() может не сработать — это нормально, продолжим диагностику
+        print("Warning: model.to(map_device) вызвал исключение:", e)
+
+    # Диагностика устройств, где лежат параметры/буферы
+    devs = print_model_devices(model)
+
+    # Выберем устройство "модели" — если их несколько, отдаём предпочтение CUDA если есть
+    if len(devs) == 0:
+        model_device = map_device
+    elif len(devs) == 1:
+        model_device = list(devs)[0]
+    else:
+        # если есть смешанные устройства — попробуем приоритет cuda, иначе первый в множестве
+        cuda_devs = [d for d in devs if 'cuda' in str(d)]
+        model_device = cuda_devs[0] if cuda_devs else list(devs)[0]
+        print("Внимание: обнаружены несколько устройств внутри state_dict(). Выбран model_device =", model_device)
+
+    # Если model_device не равен map_device — уведомим пользователя и попытаемся ещё раз загрузить с конкретным map_location
+    if str(model_device) != str(map_device):
+        print(f"Model tensors are on {model_device} but requested map_device is {map_device}.")
+        print("Попробую заново загрузить модель с map_location=model_device ...")
+        try:
+            model = torch.jit.load(str(path), map_location=str(model_device))
+            try:
+                model = model.to(model_device)
+            except Exception:
+                pass
+            devs2 = print_model_devices(model)
+            if len(devs2) == 1 and list(devs2)[0] == model_device:
+                print("Успешно перезагружено на целевое устройство.")
+        except Exception as e:
+            print("Не удалось перезагрузить модель на желаемое устройство:", e)
+            # продолжаем, но предупредим пользователя
+    return model, model_device
+
+def get_logits_from_model(model, inputs):
+    """
+    Вызов модели, допускающий возможные варианты возврата.
+    Мы предполагаем, что inputs уже находится на том же устройстве, что и модель.
+    """
+    try:
+        out = model(inputs)
+        # model может вернуть logits или (logits, kv)
+        if isinstance(out, tuple) or isinstance(out, list):
+            return out[0]
+        return out
+    except RuntimeError as e:
+        # Если ошибка связана с устройствами, добавим детальный лог
+        msg = str(e)
+        if "Expected all tensors to be on the same device" in msg or "but found at least two devices" in msg:
+            print("RuntimeError: вероятно есть mismatch устройств (cpu/cuda) внутри model. Диагностика state_dict():")
+            try:
+                print_model_devices(model)
+            except Exception:
+                pass
+            # Ребросим исключение с более понятным сообщением
+            raise RuntimeError("Device mismatch while running the JIT model. See printed diagnostics above.") from e
+        else:
+            raise
+
+# ----------------- Пример интеграции в train loop -----------------
+def train():
+    model, model_device = safe_load_jit_model(BASE_MODEL_PATH, device)
+
+    # Подготовьте датасеты здесь как вы уже делаете (замените на свой TextDataset)
+    from transformers import GPT2TokenizerFast
+    # Замените на ваш реальный TextDataset; здесь лишь заглушка
+    class DummyDataset(torch.utils.data.Dataset):
+        def __init__(self, n=1000, seq_len=TRAIN_SEQ_LEN, vocab_size=50257):
+            self.n = n
+            self.seq_len = seq_len
+            self.vocab_size = vocab_size
+        def __len__(self): return self.n
+        def __getitem__(self, i):
+            x = torch.randint(0, self.vocab_size, (self.seq_len,), dtype=torch.long)
+            y = torch.randint(0, self.vocab_size, (self.seq_len,), dtype=torch.long)
+            return x, y
+
+    train_dataset = DummyDataset(n=2000)
+    val_dataset = DummyDataset(n=200)
+
+    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)
+    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=True)
+
+    # Создаём optimizer
+    params = list(model.parameters()) if hasattr(model, 'parameters') else []
+    if len(params) == 0:
+        print("Warning: model.parameters() пуст. Убедитесь, что JIT-модель содержит параметры для оптимизации.")
+    optimizer = optim.AdamW(params, lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) if params else None
+    criterion = nn.CrossEntropyLoss()
+
+    # Инициализация GradScaler для AMP
+    scaler = GradScaler()
+
+    model.train() 
+
+    for epoch in range(1, EPOCHS + 1):
+        print(f"Эпоха {epoch}/{EPOCHS}")
+        epoch_loss = 0.0
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch} [TRAIN]", leave=False)
+        
+        batch_count = 0
+        skipped_batches = 0
+        
+        for xb, yb in pbar:
+            # === 1. ПРОВЕРКА ДАННЫХ НА NAN/INF ===
+            # Проверяем только если тип данных — float (для LongTensor проверка не нужна)
+            if torch.is_floating_point(xb) and (torch.isnan(xb).any() or torch.isinf(xb).any()):
+                 print(f"\n[E{epoch}] WARNING: NaN or Inf found in input data (xb). Skipping batch.")
+                 skipped_batches += 1
+                 continue
+            
+            # Приводим батчи к устройству модели (model_device)
+            xb = xb.to(model_device)
+            yb = yb.to(model_device)
+
+            if optimizer:
+                optimizer.zero_grad()
+
+            # === 2. AMP: Выполняем forward-pass в half-precision ===
+            with autocast():
+                logits = get_logits_from_model(model, xb)
+                
+                # У logits размер [B, seq_len, vocab] — приводим к числу классов
+                loss = criterion(logits.view(-1, logits.size(-1)), yb.view(-1))
+            # ========================================================
+            
+            # === 3. ПРОВЕРКА ЛОССА НА NAN/INF ПЕРЕД BACKWARD ===
+            # Проверяем лосс, который теперь может быть float16 или float32
+            if torch.isnan(loss) or torch.isinf(loss):
+                print(f"\n[E{epoch}] CRITICAL: Loss is NaN or Inf. Skipping backward and update.")
+                skipped_batches += 1
+                continue
+                
+            # AMP: Вычисляем градиенты, масштабируя их
+            scaler.scale(loss).backward()
+            
+            if optimizer:
+                # AMP: Сначала снимаем масштаб
+                scaler.unscale_(optimizer) 
+                
+                # Обрезка градиентов
+                torch.nn.utils.clip_grad_norm_(params, GRAD_CLIP)
+                
+                # AMP: Обновляем веса (scaler сам проверяет, не являются ли градиенты Inf/NaN)
+                scaler.step(optimizer)
+                scaler.update()
+
+            # Переводим лосс в float32 для записи и отображения
+            loss_val = loss.item()
+            epoch_loss += loss_val
+            batch_count += 1
+            
+            pbar.set_postfix({"loss": f"{loss_val:.4f}", "ppl": f"{math.exp(min(loss_val, 10)):.2f}"})
+
+        # Средняя потеря считается только по не пропущенным батчам
+        avg_loss = epoch_loss / batch_count if batch_count > 0 else float('nan')
+        print(f"Средняя потеря за эпоху: {avg_loss:.4f}")
+        
+        if skipped_batches > 0:
+            print(f"Внимание: {skipped_batches} батчей было пропущено из-за NaN/Inf в данных или лоссе.")
+
+
+if __name__ == "__main__":
+    train()

source_jit/fine_tune_native_H4_L2.py

@@ -0,0 +1,113 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+import torch
+import torch.nn as nn
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from torch.amp import autocast, GradScaler
+from tqdm import tqdm
+import math
+from pathlib import Path
+
+# Подключаем твою модель (ту же самую, что была в JIT)
+from your_model_file import JiRack_H4_L2  # ← сюда имя файла с классом модели (который я тебе дал последним)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Устройство: {device}")
+
+# Загружаем обычную модель (НЕ JIT!)
+model = JiRack_H4_L2().to(device)
+
+# Загружаем веса из JIT-конвертированной модели (они совместимы!)
+state_dict = torch.load("models/JiRack_H4_L2_V50257_D768_MSL8192_FF3072.pt", map_location=device)
+model.load_state_dict(state_dict)
+print("Веса загружены из .pt файла")
+
+# Параметры обучения — теперь можно чуть агрессивнее
+BATCH_SIZE = 12
+SEQ_LEN = 256
+EPOCHS = 10
+LR = 5e-5
+WARMUP_STEPS = 100
+
+# Твой датасет (пример с рандомом — замени на реальный)
+class DummyDataset(torch.utils.data.Dataset):
+    def __init__(self, n=10000): self.n = n
+    def __len__(self): return self.n
+    def __getitem__(self, i):
+        x = torch.randint(0, 50257, (SEQ_LEN,))
+        return x, x.roll(-1)  # next token prediction
+
+train_loader = DataLoader(DummyDataset(), batch_size=BATCH_SIZE, shuffle=True, drop_last=True)
+
+optimizer = AdamW(model.parameters(), lr=LR, weight_decay=0.01)
+scaler = GradScaler('cuda')
+criterion = nn.CrossEntropyLoss()
+
+global_step = 0
+model.train()
+
+for epoch in range(1, EPOCHS + 1):
+    total_loss = 0
+    pbar = tqdm(train_loader, desc=f"Эпоха {epoch}/{EPOCHS}")
+
+    for xb, yb in pbar:
+        global_step += 1
+        xb, yb = xb.to(device), yb.to(device)
+
+        optimizer.zero_grad()
+
+        with autocast('cuda'):
+            logits = model(xb)  # ← обычный forward, без past_kv
+            loss = criterion(logits.view(-1, logits.size(-1)), yb.view(-1))
+
+        scaler.scale(loss).backward()
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+        scaler.step(optimizer)
+        scaler.update()
+
+        # LR warmup
+        if global_step < WARMUP_STEPS:
+            lr_scale = global_step / WARMUP_STEPS
+            for pg in optimizer.param_groups:
+                pg['lr'] = LR * lr_scale
+
+        total_loss += loss.item()
+        pbar.set_postfix({"loss": f"{loss.item():.4f}", "ppl": f"{math.exp(loss.item()):.1f}"})
+
+    avg_loss = total_loss / len(train_loader)
+    print(f"Эпоха {epoch} завершена | Средний loss: {avg_loss:.4f} | Perplexity: {math.exp(avg_loss):.2f}\n")
+
+# После обучения — сохраняем и JIT-версию для инференса
+torch.save(model.state_dict(), "models/JiRack_H4_L2_finetuned.pt")
+
+# Экспорт в JIT (теперь уже обученной модели)
+class JITWrapper(nn.Module):
+    def __init__(self, m): super().__init__(); self.m = m
+    def forward(self, x): return self.m(x)
+
+dummy = torch.randint(0, 50257, (1, 256), device=device)
+traced = torch.jit.trace(JITWrapper(model.cpu().eval()), dummy)
+traced.save("models/JiRack_H4_L2_finetuned.script.pt")
+print("Обученная модель сохранена + экспортирована в JIT для инференса")

source_jit/tools_diagnostics_print_jit_constants.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+
+import sys
+import torch
+from pathlib import Path
+
+def main():
+    if len(sys.argv) < 2:
+        print("Usage: python diagnostics_print_jit_constants.py <jit_model_path>")
+        return
+    path = Path(sys.argv[1])
+    if not path.exists():
+        print("File not found:", path)
+        return
+
+    print("Loading JIT model (map_location='cpu') for safe inspection...")
+    m = torch.jit.load(str(path), map_location='cpu')
+    print("Loaded. Collecting info...\n")
+
+    print("Named parameters (name, device, shape):")
+    try:
+        for n, p in m.named_parameters():
+            print("  PARAM:", n, p.device, tuple(p.shape))
+    except Exception as e:
+        print("  (named_parameters() not available / raised):", e)
+
+    print("\nNamed buffers (name, device, shape):")
+    try:
+        for n, b in m.named_buffers():
+            print("  BUFFER:", n, b.device, tuple(b.shape))
+    except Exception as e:
+        print("  (named_buffers() not available / raised):", e)
+
+    print("\nstate_dict keys and devices:")
+    try:
+        sd = m.state_dict()
+        devices = set()
+        for k, v in sd.items():
+            try:
+                devices.add(v.device)
+                print(" ", k, v.device, tuple(v.shape))
+            except Exception:
+                print(" ", k, " - (non-tensor?)")
+        print("Devices in state_dict():", devices)
+    except Exception as e:
+        print("  state_dict() failed:", e)
+
+    print("\nAttempt to show TorchScript graph (short version). Look for prim::Constant Tensor entries:")
+    try:
+        g = m.graph
+        print(g)
+    except Exception as e:
+        print("  Could not print graph directly:", e)
+        try:
+            print("m.code():")
+            print(m.code)
+        except Exception as e2:
+            print("  Also could not print m.code():", e2)
+
+    print("\nIf you find prim::Constant values with Tensor on CPU, those likely cause device mismatch.")
+    print("Recommendation: re-create JIT on target device (see retrace_to_cuda.py).")
+
+if __name__ == "__main__":
+    main()

source_jit/tools_retrace_to_cuda.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+import argparse
+import torch
+from pathlib import Path
+import importlib
+import sys
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--jit", required=True, help="Path to existing JIT model (used to extract state_dict)")
+parser.add_argument("--out", required=True, help="Output path for new JIT model on CUDA")
+parser.add_argument("--py_module", required=False, help="Python import path for model (e.g. jirackkit.src.main.python.gpt2_jit.JiRack_H4_L2_V50257_D768_MSL8192_FF768x4)", default=None)
+parser.add_argument("--class_name", required=False, help="Name of model class in module", default=None)
+parser.add_argument("--seq_len", type=int, default=8, help="Sequence length for example input (short is fine for trace)")
+parser.add_argument("--vocab_size", type=int, default=50257, help="Vocab size for dummy input")
+parser.add_argument("--use_script", action="store_true", help="Use torch.jit.script instead of trace (requires model to be scriptable)")
+args = parser.parse_args()
+
+jit_path = Path(args.jit)
+out_path = Path(args.out)
+if not jit_path.exists():
+    print("JIT file not found:", jit_path)
+    sys.exit(1)
+
+# 1) load state_dict from existing JIT (safe: load on cpu)
+print("Loading state_dict from existing JIT (cpu)...")
+jit = torch.jit.load(str(jit_path), map_location='cpu')
+try:
+    sd = jit.state_dict()
+    print("state_dict keys:", list(sd.keys())[:10], "...")
+except Exception as e:
+    print("Failed to obtain state_dict() from JIT:", e)
+    sd = None
+
+# 2) Import python module & create model instance
+if args.py_module is None or args.class_name is None:
+    print("ERROR: You must provide --py_module and --class_name to reconstruct the Python model.")
+    print("Example: --py_module jirackkit.src.main.python.gpt2_jit.JiRack_H4_L2_V50257_D768_MSL8192_FF768x4 --class_name GPTPyTorch")
+    sys.exit(1)
+
+print("Importing Python model:", args.py_module, args.class_name)
+module = importlib.import_module(args.py_module)
+ModelClass = getattr(module, args.class_name)
+
+# NOTE: Provide the correct constructor args for your model here if needed.
+MODEL_KWARGS = {}  # <-- EDIT if your model constructor requires arguments
+
+print("Instantiating Python model...")
+model = ModelClass(**MODEL_KWARGS)
+
+# 3) load weights if available
+if sd is not None:
+    try:
+        model.load_state_dict(sd)
+        print("Weights loaded into Python model from JIT.state_dict().")
+    except Exception as e:
+        print("Failed to load state_dict into Python model:", e)
+        print("You may need to adapt keys or load partial weights. Exiting.")
+        sys.exit(1)
+
+# 4) move to cuda
+if not torch.cuda.is_available():
+    print("CUDA not available on this machine. Aborting.")
+    sys.exit(1)
+device = torch.device('cuda:0')
+model.to(device)
+model.eval()
+
+# 5) prepare example input on CUDA (batch=1)
+seq_len = args.seq_len
+vocab = args.vocab_size
+example_input = torch.randint(0, vocab, (1, seq_len), dtype=torch.long, device=device)
+
+# 6) trace or script
+print("Tracing/script-model on CUDA. This will produce a JIT module whose constants are on CUDA.")
+if args.use_script:
+    print("Using torch.jit.script...")
+    scripted = torch.jit.script(model)
+else:
+    print("Using torch.jit.trace with example input of shape", example_input.shape)
+    scripted = torch.jit.trace(model, example_input)
+
+# 7) save
+out_path.parent.mkdir(parents=True, exist_ok=True)
+scripted.save(str(out_path))
+print("Saved new JIT (CUDA) model to:", out_path)
+print("Done. Replace your old model file with this one (keep backup).")