Upload train_v3.py with huggingface_hub

train_v3.py

@@ -0,0 +1,288 @@
+#!/usr/bin/env python3
+"""Parallax-Chess v3: Predict move as single index (from_sq * 64 + to_sq).
+Fixes the v2 bug where only from-square was supervised.
+"""
+import sys, math, time, json, random
+from pathlib import Path
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import chess
+
+ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(ROOT / "scripts"))
+sys.path.insert(0, str(ROOT / "src"))
+sys.path.insert(0, str(ROOT.parent))
+from centauri.models.base import SmallLMConfig
+from centauri.models.torch.small_lm import SmallLM
+from chess_encoding import encode_board, VOCAB_SIZE
+
+random.seed(42)
+
+# Move index: from_sq * 64 + to_sq for non-promo, then promo moves
+# Total: 64*64 = 4096 normal + 4*8*8 = 256 promo = 4352 total
+MOVE_INDEX_SIZE = 4352  # 64*64 + 4*64 (promotions)
+
+
+def move_to_index(move):
+    """Convert chess.Move to index 0..4351."""
+    if move.promotion:
+        promo_map = {chess.QUEEN: 0, chess.ROOK: 1, chess.BISHOP: 2, chess.KNIGHT: 3}
+        return 4096 + promo_map[move.promotion] * 64 + move.to_square
+    return move.from_square * 64 + move.to_square
+
+
+def index_to_move(idx, board):
+    """Convert index back to chess.Move, only if legal."""
+    if idx >= 4096:
+        promo_offset = idx - 4096
+        promo_type = promo_offset // 64
+        to_sq = promo_offset % 64
+        promo_map = {0: chess.QUEEN, 1: chess.ROOK, 2: chess.BISHOP, 3: chess.KNIGHT}
+        promotion = promo_map[promo_type]
+        # Find from_square (must be a pawn on correct file)
+        for from_sq in chess.SQUARES:
+            piece = board.piece_at(from_sq)
+            if piece and piece.piece_type == chess.PAWN and piece.color == board.turn:
+                move = chess.Move(from_sq, to_sq, promotion=promotion)
+                if move in board.legal_moves:
+                    return move
+        return None
+    else:
+        from_sq = idx // 64
+        to_sq = idx % 64
+        move = chess.Move(from_sq, to_sq)
+        if move in board.legal_moves:
+            return move
+        return None
+
+
+class ChessDataset(Dataset):
+    """Load SF-evaluated positions."""
+
+    def __init__(self, data_paths):
+        self.samples = []
+        for path in data_paths:
+            path = Path(path)
+            if not path.exists():
+                continue
+            with open(path, "r", encoding="utf-8") as f:
+                for line in f:
+                    try:
+                        item = json.loads(line)
+                        board = chess.Board(item["fen"])
+                        move_key = "best_move" if "best_move" in item else "move"
+                        move = chess.Move.from_uci(item[move_key])
+                        if move not in board.legal_moves:
+                            continue
+
+                        board_tokens = encode_board(board)
+                        move_idx = move_to_index(move)
+
+                        if "eval_cp" in item:
+                            value = max(-1.0, min(1.0, item["eval_cp"] / 1000.0))
+                        elif "result" in item:
+                            r = item["result"]
+                            value = 1.0 if r == "1-0" else -1.0 if r == "0-1" else 0.0
+                        else:
+                            value = 0.0
+
+                        self.samples.append((board_tokens, move_idx, value))
+
+                        if random.random() < 0.5:
+                            flipped = flip_board(board_tokens)
+                            flipped_move = flip_move_index(move_idx)
+                            self.samples.append((flipped, flipped_move, value))
+                    except:
+                        pass
+        print("Loaded %d samples" % len(self.samples))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, i):
+        board_tokens, move_idx, value = self.samples[i]
+        return (
+            torch.tensor(board_tokens, dtype=torch.long),
+            torch.tensor(move_idx, dtype=torch.long),
+            torch.tensor(value, dtype=torch.float),
+        )
+
+
+def flip_board(board_tokens):
+    flipped = list(board_tokens)
+    for r in range(8):
+        for f in range(8):
+            flipped[r * 8 + (7 - f)] = board_tokens[r * 8 + f]
+    return flipped
+
+
+def flip_sq(sq):
+    """Mirror a square horizontally: file f -> 7-f."""
+    rank = sq // 8
+    file = sq % 8
+    return rank * 8 + (7 - file)
+
+
+def flip_move_index(idx):
+    """Flip move index when board is mirrored horizontally."""
+    if idx >= 4096:
+        promo_offset = idx - 4096
+        promo_type = promo_offset // 64
+        to_sq = promo_offset % 64
+        return 4096 + promo_type * 64 + flip_sq(to_sq)
+    from_sq = idx // 64
+    to_sq = idx % 64
+    return flip_sq(from_sq) * 64 + flip_sq(to_sq)
+
+
+class ParallaxChessV3(nn.Module):
+    """Board -> move index (4352 classes) + value."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.board_encoder = nn.Embedding(VOCAB_SIZE, cfg.d_model)
+        from centauri.models.torch.small_lm import SmallLM
+        self.backbone = SmallLM(cfg)
+        self.policy_head = nn.Linear(cfg.d_model, MOVE_INDEX_SIZE)
+        self.value_head = nn.Sequential(
+            nn.Linear(cfg.d_model, 256), nn.ReLU(), nn.Dropout(0.1),
+            nn.Linear(256, 1), nn.Tanh())
+
+    def forward(self, board_tokens):
+        x = self.board_encoder(board_tokens)
+        freq, _ = __import__('centauri.models.torch.small_lm', fromlist=['precompute_rope']).precompute_rope(
+            self.cfg, x.device, x.dtype)
+        for layer in self.backbone.layers:
+            x, _ = layer(x, freq, None)
+        pooled = x.mean(dim=1)
+        return {"logits": self.policy_head(pooled), "value": self.value_head(pooled).squeeze(-1)}
+
+    def predict_move(self, board, device=None):
+        if device is None:
+            device = next(self.parameters()).device
+        board_tokens = encode_board(board)
+        inp = torch.tensor([board_tokens], dtype=torch.long).to(device)
+        with torch.no_grad():
+            out = self(inp)
+            probs = torch.softmax(out["logits"], dim=-1)
+
+        # Score only legal moves
+        best_move, best_score = None, -float("inf")
+        for move in board.legal_moves:
+            idx = move_to_index(move)
+            score = probs[0, idx].item()
+            if score > best_score:
+                best_score = score
+                best_move = move
+        return best_move
+
+    def evaluate(self, board, device=None):
+        if device is None:
+            device = next(self.parameters()).device
+        board_tokens = encode_board(board)
+        inp = torch.tensor([board_tokens], dtype=torch.long).to(device)
+        with torch.no_grad():
+            return self(inp)["value"].item() * 1000
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--max_steps", type=int, default=80000)
+    parser.add_argument("--save_dir", default=str(ROOT / "checkpoints" / "parallax_chess_v3"))
+    args = parser.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Device:", device)
+
+    cfg_dict = {
+        "vocab_size": VOCAB_SIZE, "d_model": 512, "n_heads": 8, "n_kv_heads": 4,
+        "n_layers": 8, "intermediate_size": 2048, "max_seq_len": 128,
+        "norm_type": "rms", "rope_type": "neox", "n_experts": 0,
+        "n_loops": 1, "loop_mode": "per_layer",
+    }
+    cfg = SmallLMConfig(**cfg_dict)
+    model = ParallaxChessV3(cfg).to(device)
+    n_params = sum(p.numel() for p in model.parameters())
+    print("Params: %d (%.1fM)" % (n_params, n_params / 1e6))
+
+    data_paths = [
+        ROOT / "data" / "chess_train_sf.jsonl",
+        ROOT / "data" / "chess_train_large.jsonl",
+    ]
+    ds = ChessDataset(data_paths)
+    loader = DataLoader(ds, batch_size=args.batch_size, shuffle=True, num_workers=0, pin_memory=True)
+
+    save_dir = Path(args.save_dir)
+    save_dir.mkdir(parents=True, exist_ok=True)
+
+    step = 0
+    for f in sorted(save_dir.glob("step_*.pt")):
+        try:
+            ckpt = torch.load(str(f), map_location=device)
+            model.load_state_dict(ckpt["model"])
+            step = ckpt.get("step", 0)
+            print("Resumed from %s at step %d" % (f.name, step))
+            break
+        except:
+            pass
+
+    total_steps = min(args.max_steps, len(ds) // args.batch_size * 5)
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01, betas=(0.9, 0.98))
+
+    def lr_schedule(s):
+        warmup = 2000
+        if s < warmup:
+            return s / warmup
+        progress = (s - warmup) / max(1, total_steps - warmup)
+        return max(0.05, 0.5 * (1.0 + math.cos(math.pi * progress)))
+
+    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_schedule)
+
+    model.train()
+    t0 = time.time()
+
+    for epoch in range(999):
+        for board_tokens, move_indices, values in loader:
+            board_tokens = board_tokens.to(device)
+            move_indices = move_indices.to(device)
+            values = values.to(device)
+
+            out = model(board_tokens)
+
+            policy_loss = F.cross_entropy(out["logits"], move_indices)
+            value_loss = F.mse_loss(out["value"], values)
+            loss = policy_loss + 1.0 * value_loss
+
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            opt.step()
+            opt.zero_grad()
+            sched.step()
+
+            if step % 100 == 0:
+                lr = sched.get_last_lr()[0]
+                print("Step %6d | P:%.4f V:%.4f | Loss:%.4f | LR %.1e" % (
+                    step, policy_loss.item(), value_loss.item(), loss.item(), lr))
+
+            if step > 0 and step % 10000 == 0:
+                ckpt = save_dir / ("step_%d.pt" % step)
+                torch.save({"model": model.state_dict(), "step": step, "config": cfg_dict, "n_params": n_params}, str(ckpt))
+                print("Saved %s" % ckpt.name)
+
+            step += 1
+            if step >= total_steps:
+                break
+
+    final = save_dir / "final.pt"
+    torch.save({"model": model.state_dict(), "step": step, "config": cfg_dict, "n_params": n_params}, str(final))
+    print("Done! %s (%d steps)" % (final, step))
+
+
+if __name__ == "__main__":
+    main()