Upload modeling_kiyengine.py with huggingface_hub

modeling_kiyengine.py

@@ -0,0 +1,294 @@
+
+# === Imports ===
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import chess
+import chess.pgn
+import os
+import random
+import pickle
+import time
+import glob
+from typing import Dict, List, Tuple
+from tqdm import tqdm
+from safetensors.torch import save_file
+from torch.utils.data import Dataset, DataLoader
+from torch.amp import GradScaler, autocast
+
+# === Configuration (P100 Optimized & FIXED) ===
+CONFIG = {
+    'model': {
+        'd_model': 384, 'n_layers': 4, 'n_experts': 8, 'top_k': 2, 'd_state': 16,
+        'd_conv': 4, 'expansion_factor': 2, 'vocab_size': 768,
+    },
+    'training': {
+        'batch_size': 4096,      
+        'learning_rate': 4.0e-4, 
+        'epochs': 10, 
+        'noise_sigma': 0.01,
+        'save_every_mins': 15,   
+        'keep_checkpoints': 2,
+        # --- [FIX HERE] Trả lại các trọng số đã bị thất lạc ---
+        'policy_weight': 1.0,
+        'value_weight': 1.0,
+        'aux_loss_lambda': 0.01,
+        # -----------------------------------------------------
+    },
+    'paths': {
+        'train_data_path': "/kaggle/working/train_data.pgn",
+        'save_path': "./snapshots", 
+        'model_save_name': "model.safetensors",
+    },
+}
+
+# === Helper: Data Prefetcher ===
+class DataPrefetcher:
+    def __init__(self, loader, device):
+        self.loader = iter(loader)
+        self.device = device
+        self.stream = torch.cuda.Stream()
+        self.preload()
+
+    def preload(self):
+        try:
+            self.next_batch = next(self.loader)
+        except StopIteration:
+            self.next_batch = None
+            return
+
+        with torch.cuda.stream(self.stream):
+            self.next_batch = [x.to(self.device, non_blocking=True) for x in self.next_batch]
+
+    def next(self):
+        torch.cuda.current_stream().wait_stream(self.stream)
+        batch = self.next_batch
+        self.preload()
+        return batch
+
+# === Helper: Rolling Checkpoint Manager ===
+def manage_checkpoints(save_dir, keep_n=2):
+    files = glob.glob(os.path.join(save_dir, "checkpoint_*.safetensors"))
+    files.sort(key=os.path.getmtime)
+    while len(files) > keep_n:
+        oldest_file = files.pop(0)
+        try:
+            os.remove(oldest_file)
+            print(f"🗑️ Cleaned up old checkpoint: {oldest_file}")
+        except OSError as e:
+            print(f"⚠️ Error deleting file {oldest_file}: {e}")
+
+# === Model Architecture (Mamba + MoE) ===
+class GaussianNoise(nn.Module):
+    def __init__(self, sigma: float = 0.01): super().__init__(); self.sigma = sigma
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.training and self.sigma != 0: return x + torch.randn_like(x) * self.sigma
+        return x
+
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-5): 
+        super().__init__(); self.eps = eps; self.weight = nn.Parameter(torch.ones(d_model))
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        norm = x.norm(2, dim=-1, keepdim=True) * (x.shape[-1] ** -0.5)
+        return x / (norm + self.eps) * self.weight
+
+class MambaBlock(nn.Module):
+    def __init__(self, config: Dict): 
+        super().__init__()
+        d_model, d_state, d_conv, exp_factor = config['d_model'], config['d_state'], config['d_conv'], config['expansion_factor']
+        d_inner = d_model * exp_factor
+        self.in_proj = nn.Linear(d_model, 2 * d_inner, bias=False)
+        self.conv1d = nn.Conv1d(in_channels=d_inner, out_channels=d_inner, kernel_size=d_conv, bias=True, groups=d_inner, padding=d_conv - 1)
+        self.x_proj = nn.Linear(d_inner, d_inner + 2 * d_state, bias=False)
+        self.dt_proj = nn.Linear(d_inner, d_inner, bias=True)
+        self.A_log = nn.Parameter(torch.randn(d_inner, d_state)); self.D = nn.Parameter(torch.ones(d_inner))
+        self.out_proj = nn.Linear(d_inner, d_model, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, L, C = x.shape; xz = self.in_proj(x); x_inner, z = xz.chunk(2, dim=-1)
+        x_conv = self.conv1d(x_inner.transpose(1, 2))[:, :, :L].transpose(1, 2); x_activated = F.silu(x_conv)
+        y = x_activated * self.D.unsqueeze(0); y = y * F.silu(z)
+        return self.out_proj(y)
+
+class MoELayer(nn.Module):
+    def __init__(self, config: Dict): 
+        super().__init__(); self.n_experts, self.top_k = config['n_experts'], config['top_k']
+        self.router = nn.Linear(config['d_model'], self.n_experts)
+        self.experts = nn.ModuleList([MambaBlock(config) for _ in range(self.n_experts)])
+
+    def forward(self, x: torch.Tensor) -> (torch.Tensor, torch.Tensor):
+        B, L, C = x.shape; x_flat = x.view(-1, C); router_logits = self.router(x_flat)
+        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+        top_k_weights, top_k_indices = torch.topk(routing_weights, self.top_k, dim=-1)
+        top_k_weights /= top_k_weights.sum(dim=-1, keepdim=True)
+        expert_mask = F.one_hot(top_k_indices, self.n_experts).sum(dim=1); expert_load = expert_mask.float().mean(dim=0)
+        aux_loss = (expert_load * expert_load).sum()
+        final_output = torch.zeros_like(x_flat)
+        for i in range(self.top_k):
+            expert_idx = top_k_indices[:, i]; weight = top_k_weights[:, i].unsqueeze(-1)
+            for j in range(self.n_experts):
+                mask = expert_idx == j
+                if mask.any(): final_output[mask] += (self.experts[j](x_flat[mask].unsqueeze(1)).squeeze(1) * weight[mask])
+        return final_output.view(B, L, C), aux_loss
+
+class KiyEngineV3(nn.Module):
+    def __init__(self, config: Dict): 
+        super().__init__(); self.config = config
+        self.embedding = nn.Embedding(config['vocab_size'], config['d_model'])
+        self.noise = GaussianNoise(sigma=config.get('training', {}).get('noise_sigma', 0.0))
+        self.layers = nn.ModuleList([MoELayer(config) for _ in range(config['n_layers'])])
+        self.norm = RMSNorm(config['d_model'])
+        self.policy_head = nn.Linear(config['d_model'], config['vocab_size'], bias=False)
+        self.value_head = nn.Sequential(nn.Linear(config['d_model'], 128), nn.ReLU(), nn.Linear(128, 1))
+
+    def forward(self, input_ids: torch.Tensor) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+        x = self.noise(self.embedding(input_ids)); total_aux_loss = 0.0
+        for layer in self.layers: x = x + layer(self.norm(x))[0]; total_aux_loss += layer(self.norm(x))[1]
+        x = self.norm(x); last_token_state = x[:, -1, :]
+        policy_logits = self.policy_head(last_token_state); value = torch.tanh(self.value_head(last_token_state))
+        return policy_logits, value, total_aux_loss / self.config['n_layers']
+
+# === Data Pipeline (Header Only + Robust) ===
+def move_to_token(move, board):
+    piece = board.piece_at(move.from_square)
+    if piece is None: return 0 
+    piece_idx = move.promotion - 1 if move.promotion else piece.piece_type - 1
+    if piece.color == chess.BLACK: piece_idx += 6
+    return piece_idx * 64 + move.to_square
+
+class ChessDataset(Dataset):
+    def __init__(self, pgn_file_path, context_length=16): 
+        self.pgn_file_path = pgn_file_path
+        self.context_length = context_length
+        self.games = self._index_games(pgn_file_path)
+
+    def _index_games(self, pgn_file_path):
+        index_path = pgn_file_path + ".index.pkl"
+        if os.path.exists(index_path):
+             print(f"🚀 Loading cached index from {index_path}...")
+             with open(index_path, "rb") as f: return pickle.load(f)
+
+        print(f"⚡ Turbo Indexing {pgn_file_path} (Header Only Mode)...")
+        offsets = []
+        count = 0
+        with open(pgn_file_path) as pgn:
+            while True:
+                offset = pgn.tell()
+                headers = chess.pgn.read_headers(pgn)
+                if headers is None: break
+                res = headers.get("Result", "*")
+                val = 0.0
+                if res == "1-0": val = 1.0
+                elif res == "0-1": val = -1.0
+                elif res == "1/2-1/2": val = 0.0
+                else: continue 
+                offsets.append((offset, val))
+                count += 1
+                if count % 50000 == 0: print(f"Indexed {count} games...", end='\r')
+        
+        print(f"\n✅ Done! Found {len(offsets)} valid games.")
+        with open(index_path, "wb") as f: pickle.dump(offsets, f)
+        return offsets
+
+    def __len__(self): return len(self.games)
+
+    def __getitem__(self, idx):
+        offset, value = self.games[idx]
+        try:
+            with open(self.pgn_file_path) as f:
+                f.seek(offset)
+                game = chess.pgn.read_game(f)
+                if game is None or game.errors: return torch.zeros(self.context_length, dtype=torch.long), torch.tensor(0, dtype=torch.long), torch.tensor([0.0])
+                moves = list(game.mainline_moves())
+                if len(moves) <= self.context_length: return torch.zeros(self.context_length, dtype=torch.long), torch.tensor(0, dtype=torch.long), torch.tensor([0.0])
+
+                start_ply = random.randint(0, len(moves) - self.context_length - 1)
+                move_history = moves[start_ply : start_ply + self.context_length]
+                target_move = moves[start_ply + self.context_length]
+                
+                board = chess.Board()
+                for i in range(start_ply): board.push(moves[i])
+                temp_board = board.copy()
+                seq = []
+                for move in move_history: 
+                    seq.append(move_to_token(move, temp_board))
+                    temp_board.push(move)
+                target_token = move_to_token(target_move, temp_board)
+                return torch.tensor(seq), torch.tensor(target_token), torch.tensor([value])
+        except Exception:
+            return torch.zeros(self.context_length, dtype=torch.long), torch.tensor(0, dtype=torch.long), torch.tensor([0.0])
+
+# === Training Loop (Single GPU Optimized) ===
+def train_main():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"🔥 Hardware: {torch.cuda.get_device_name(0)}")
+
+    os.makedirs(CONFIG['paths']['save_path'], exist_ok=True)
+    
+    model = KiyEngineV3(CONFIG['model']).to(device)
+    dataset = ChessDataset(CONFIG['paths']['train_data_path'])
+    
+    dataloader = DataLoader(dataset, batch_size=CONFIG['training']['batch_size'], 
+                            shuffle=True, num_workers=os.cpu_count(), pin_memory=True)
+    
+    optimizer = optim.Adam(model.parameters(), lr=CONFIG['training']['learning_rate'])
+    scaler = GradScaler('cuda')
+
+    print("🚀 Starting P100 Turbo Training...")
+    
+    last_save_time = time.time()
+    
+    for epoch in range(CONFIG['training']['epochs']):
+        prefetcher = DataPrefetcher(dataloader, device)
+        batch = prefetcher.next()
+        
+        pbar = tqdm(total=len(dataloader), desc=f"Epoch {epoch+1}")
+        
+        batch_idx = 0
+        while batch is not None:
+            input_seq, policy_target, value_target = batch
+            
+            optimizer.zero_grad()
+            with autocast('cuda'): 
+                policy_logits, value_pred, aux_loss = model(input_seq)
+                policy_loss = F.cross_entropy(policy_logits, policy_target)
+                value_loss = F.mse_loss(value_pred.squeeze(), value_target.squeeze())
+                # --- Hàng về rồi đây ---
+                loss = CONFIG['training']['policy_weight'] * policy_loss + CONFIG['training']['value_weight'] * value_loss + CONFIG['training']['aux_loss_lambda'] * aux_loss
+            
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            
+            if (time.time() - last_save_time) > (CONFIG['training']['save_every_mins'] * 60):
+                checkpoint_name = f"checkpoint_ep{epoch+1}_step{batch_idx}.safetensors"
+                save_path = os.path.join(CONFIG['paths']['save_path'], checkpoint_name)
+                
+                model_to_save = model
+                tensors = {name: param for name, param in model_to_save.state_dict().items()}
+                save_file(tensors, save_path)
+                
+                print(f"\n💾 Auto-saved: {checkpoint_name}")
+                manage_checkpoints(CONFIG['paths']['save_path'], keep_n=CONFIG['training']['keep_checkpoints'])
+                last_save_time = time.time()
+
+            if batch_idx % 100 == 0:
+                 with open("training_progress.log", "a") as f:
+                    f.write(f"Epoch {epoch+1} | Batch {batch_idx} | Loss: {loss.item():.4f}\n")
+            
+            pbar.set_postfix({"Loss": f"{loss.item():.4f}"})
+            pbar.update(1)
+            
+            batch = prefetcher.next()
+            batch_idx += 1
+        
+        pbar.close()
+
+    final_path = os.path.join(CONFIG['paths']['save_path'], CONFIG['paths']['model_save_name'])
+    tensors = {name: param for name, param in model.state_dict().items()}
+    save_file(tensors, final_path)
+    print(f"🏁 Model saved to {final_path}")
+
+if __name__ == "__main__":
+    train_main()