Upload folder using huggingface_hub

chess_io.py

@@ -1,15 +1,4 @@
-"""Move encoding shared by data pipeline, model, search, and the UCI engine.
-
-Move representation: every move is (from_square, to_square, promo_idx).
-promo_idx in {0: none, 1: q, 2: r, 3: b, 4: n}. The from/to squares alone give a
-4096-way index (from_sq * 64 + to_sq) that covers every geometrically possible
-square pair -- no separate move-vocabulary table needed, and it trivially
-covers every legal move in any position since legality is enforced by masking,
-not by the encoding.
-"""
-
 from __future__ import annotations
-
 import chess
 
 PROMO_PIECES = [None, chess.QUEEN, chess.ROOK, chess.BISHOP, chess.KNIGHT]
@@ -19,10 +8,9 @@ NUM_PROMO = len(PROMO_PIECES)
 
 
 def move_to_ids(move: chess.Move) -> tuple[int, int]:
-    """Return (from_to_id in [0,4096), promo_id in [0,5))."""
     from_to_id = move.from_square * 64 + move.to_square
     promo_id = PROMO_INDEX[move.promotion]
-    return from_to_id, promo_id
+    return (from_to_id, promo_id)
 
 
 def ids_to_move(from_to_id: int, promo_id: int) -> chess.Move:
@@ -35,10 +23,5 @@ def legal_move_ids(board: chess.Board) -> list[tuple[int, int]]:
 
 
 def find_legal_move(board: chess.Board, from_to_id: int, promo_id: int) -> chess.Move | None:
-    """Match a predicted (from_to_id, promo_id) against the board's actual legal moves.
-
-    Needed because promo_id alone doesn't disambiguate underpromotion capture direction
-    edge cases are already fully specified by from/to, so this is just a legality check.
-    """
     candidate = ids_to_move(from_to_id, promo_id)
     return candidate if candidate in board.legal_moves else None

engine_uci.py

@@ -1,26 +1,16 @@
-"""UCI protocol loop for the ChessModel entry.
-
-Safety-net design: this runs and plays LEGAL moves even with no trained checkpoint
-present (falls back to a policy-free shallow search / random legal move). The
-trained ChessMamba net is loaded opportunistically -- if the checkpoint file isn't
-there yet, the engine still works, just weaker. This guarantees there is always
-something submittable regardless of how training goes under the deadline.
-"""
-
 from __future__ import annotations
-
 import os
 import random
 import sys
-
 import chess
-
 from search import Searcher
 
 ENGINE_NAME = "ChessMamba"
 ENGINE_AUTHOR = "CatGirlHtfi"
-DEFAULT_CKPT = os.environ.get("CHESSMAMBA_CKPT", os.path.join(os.path.dirname(__file__), "ckpt", "model.pt"))
-POLICY_ONLY = os.environ.get("CHESSMAMBA_POLICY_ONLY", "") == "1"  # diagnostic: skip negamax, play raw top policy move
+DEFAULT_CKPT = os.environ.get(
+    "CHESSMAMBA_CKPT", os.path.join(os.path.dirname(__file__), "ckpt", "model.pt")
+)
+POLICY_ONLY = os.environ.get("CHESSMAMBA_POLICY_ONLY", "") == "1"
 
 
 def send(msg: str) -> None:
@@ -31,14 +21,12 @@ def send(msg: str) -> None:
 def main() -> None:
     board = chess.Board()
     searcher = Searcher(checkpoint_path=DEFAULT_CKPT, policy_only=POLICY_ONLY)
-
     for line in sys.stdin:
         line = line.strip()
         if not line:
             continue
         tokens = line.split()
         cmd = tokens[0]
-
         if cmd == "uci":
             send(f"id name {ENGINE_NAME}")
             send(f"id author {ENGINE_AUTHOR}")
@@ -64,7 +52,6 @@ def main() -> None:
             break
         elif cmd == "stop":
             pass
-        # unknown commands are ignored, per UCI spec
 
 
 def _handle_position(board: chess.Board, tokens: list[str]) -> None:
@@ -79,7 +66,6 @@ def _handle_position(board: chess.Board, tokens: list[str]) -> None:
         idx = moves_idx
     else:
         return
-
     if idx < len(tokens) and tokens[idx] == "moves":
         for uci_move in tokens[idx + 1 :]:
             board.push_uci(uci_move)

model.py

@@ -1,17 +1,5 @@
-"""ChessMamba: a from-scratch selective-state-space (S6) model over move-history
-sequences, with a policy head (masked to legal moves) and a value head.
-
-Unlike NNUE (static sparse-feature eval + alpha-beta) or AlphaZero/Lc0 (CNN/ResNet
-over board planes + MCTS), this treats a chess game as a causal sequence of moves
-and scans it with a Mamba-style selective SSM, matching how "Chess Llama" and
-DeepMind's search-free transformer frame the problem but swapping attention for a
-recurrent scan. See PLAN.md for the full rationale.
-"""
-
 from __future__ import annotations
-
 import math
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -19,11 +7,12 @@ from torch.utils.checkpoint import checkpoint
 
 NUM_FROM_TO = 4096
 NUM_PROMO = 5
-MAX_PLIES = 96  # window of recent half-moves fed to the model
+MAX_PLIES = 96
 
 
 class RMSNorm(nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-5):
+
+    def __init__(self, dim: int, eps: float = 1e-05):
         super().__init__()
         self.eps = eps
         self.weight = nn.Parameter(torch.ones(dim))
@@ -34,22 +23,11 @@ class RMSNorm(nn.Module):
 
 
 def parallel_scan(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-    """Hillis-Steele (doubling) associative scan for the affine recurrence
-    h_t = a_t * h_{t-1} + b_t, h_{-1} = 0, computed along dim=1 (time).
-
-    a, b: (B, L, ...). Returns h: same shape, h[:, t] = the recurrence's value at t.
-
-    This is the same mathematical trick real Mamba's CUDA parallel-scan kernel uses
-    to make a sequential recurrence run in O(log L) dependent steps instead of O(L):
-    the affine maps x -> a*x + b compose associatively, (a1,b1) then (a2,b2) gives
-    (a2*a1, a2*b1+b2), so prefix-composing all the way to t gives h_t directly (since
-    composing the identity initial state h_{-1}=0 just picks out the b component).
-    """
     L = a.shape[1]
     d = 1
     while d < L:
-        a_prev, b_prev = a[:, :-d], b[:, :-d]
-        a_cur, b_cur = a[:, d:], b[:, d:]
+        a_prev, b_prev = (a[:, :-d], b[:, :-d])
+        a_cur, b_cur = (a[:, d:], b[:, d:])
         new_a = a_cur * a_prev
         new_b = a_cur * b_prev + b_cur
         a = torch.cat([a[:, :d], new_a], dim=1)
@@ -59,17 +37,6 @@ def parallel_scan(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
 
 
 class S6Block(nn.Module):
-    """Selective-scan SSM block (the core of Mamba), parallel-scan implementation.
-
-    The recurrence h_t = A_bar_t * h_{t-1} + Bx_t is sequential by nature, but it's
-    also an affine recurrence, so it can be computed via an O(log L) doubling scan
-    (see `parallel_scan`) instead of an O(L) Python loop. Measured impact: the
-    naive per-timestep loop (depth=10, ~96 timesteps, plus the recompute pass from
-    gradient checkpointing) bottlenecked training at ~11s/step on an A100 -- almost
-    entirely fixed per-launch overhead on a chain of ~1920 tiny sequential kernel
-    launches per step that batch size can't parallelize away, since batching only
-    parallelizes *within* a timestep, not across the dependency chain.
-    """
 
     def __init__(self, dim: int, state_dim: int = 16, expand: int = 2):
         super().__init__()
@@ -77,79 +44,56 @@ class S6Block(nn.Module):
         self.dim = dim
         self.inner_dim = inner_dim
         self.state_dim = state_dim
-
-        self.in_proj = nn.Linear(dim, inner_dim * 2, bias=False)  # -> (x, gate)
-        # input-dependent SSM parameters, projected from x
-        self.x_proj = nn.Linear(inner_dim, state_dim * 2 + inner_dim, bias=False)  # -> (B, C, delta_raw)
+        self.in_proj = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.x_proj = nn.Linear(inner_dim, state_dim * 2 + inner_dim, bias=False)
         self.dt_bias = nn.Parameter(torch.zeros(inner_dim))
-
-        # A: per-channel state matrix, log-parameterized for stability (stays negative after -exp)
         A = torch.arange(1, state_dim + 1, dtype=torch.float32).unsqueeze(0).repeat(inner_dim, 1)
         self.A_log = nn.Parameter(torch.log(A))
         self.D = nn.Parameter(torch.ones(inner_dim))
-
         self.out_proj = nn.Linear(inner_dim, dim, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x: (B, L, dim)
         B_, L, _ = x.shape
         xz = self.in_proj(x)
-        x_in, gate = xz.chunk(2, dim=-1)  # each (B, L, inner_dim)
+        x_in, gate = xz.chunk(2, dim=-1)
         x_in = F.silu(x_in)
-
-        x_dbl = self.x_proj(x_in)  # (B, L, state_dim*2 + inner_dim)
+        x_dbl = self.x_proj(x_in)
         Bmat, Cmat, delta_raw = torch.split(
             x_dbl, [self.state_dim, self.state_dim, self.inner_dim], dim=-1
         )
-        delta = F.softplus(delta_raw + self.dt_bias)  # (B, L, inner_dim), > 0
-
-        A = -torch.exp(self.A_log)  # (inner_dim, state_dim), negative
-
-        # discretize for every timestep at once: A_bar_t = exp(delta_t * A), Bx_t = delta_t * B_t * x_t
-        A_bar = torch.exp(delta.unsqueeze(-1) * A.view(1, 1, self.inner_dim, self.state_dim))  # (B,L,inner_dim,state_dim)
-        Bx = (delta * x_in).unsqueeze(-1) * Bmat.unsqueeze(2)  # (B, L, inner_dim, state_dim)
-
-        h = parallel_scan(A_bar, Bx)  # (B, L, inner_dim, state_dim), h[:, t] = hidden state after step t
-
-        y = (h * Cmat.unsqueeze(2)).sum(-1) + self.D * x_in  # (B, L, inner_dim)
+        delta = F.softplus(delta_raw + self.dt_bias)
+        A = -torch.exp(self.A_log)
+        A_bar = torch.exp(delta.unsqueeze(-1) * A.view(1, 1, self.inner_dim, self.state_dim))
+        Bx = (delta * x_in).unsqueeze(-1) * Bmat.unsqueeze(2)
+        h = parallel_scan(A_bar, Bx)
+        y = (h * Cmat.unsqueeze(2)).sum(-1) + self.D * x_in
         y = y * F.silu(gate)
         return self.out_proj(y)
 
-    def step(self, x_t: torch.Tensor, h_prev: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]:
-        """Single-timestep update for live/incremental play: given the token at the
-        current step and the carried recurrent state, returns (output, new_state) --
-        O(1) per move instead of re-scanning the whole move history from scratch.
-        This is the exact per-step formula `parallel_scan` is mathematically
-        equivalent to (verified against it in tests), just applied once instead of
-        via the doubling trick, since there's nothing to parallelize over a single step.
-
-        x_t: (B, dim). h_prev: (B, inner_dim, state_dim) or None (treated as zero,
-        i.e. the initial state before any moves).
-        """
+    def step(
+        self, x_t: torch.Tensor, h_prev: torch.Tensor | None
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         xz = self.in_proj(x_t)
-        x_in, gate = xz.chunk(2, dim=-1)  # (B, inner_dim) each
+        x_in, gate = xz.chunk(2, dim=-1)
         x_in = F.silu(x_in)
-
         x_dbl = self.x_proj(x_in)
         Bmat, Cmat, delta_raw = torch.split(
             x_dbl, [self.state_dim, self.state_dim, self.inner_dim], dim=-1
         )
-        delta = F.softplus(delta_raw + self.dt_bias)  # (B, inner_dim)
-
-        A = -torch.exp(self.A_log)  # (inner_dim, state_dim)
-        A_bar = torch.exp(delta.unsqueeze(-1) * A.unsqueeze(0))  # (B, inner_dim, state_dim)
-        Bx = (delta * x_in).unsqueeze(-1) * Bmat.unsqueeze(1)  # (B, inner_dim, state_dim)
-
+        delta = F.softplus(delta_raw + self.dt_bias)
+        A = -torch.exp(self.A_log)
+        A_bar = torch.exp(delta.unsqueeze(-1) * A.unsqueeze(0))
+        Bx = (delta * x_in).unsqueeze(-1) * Bmat.unsqueeze(1)
         if h_prev is None:
             h_prev = x_t.new_zeros(x_t.shape[0], self.inner_dim, self.state_dim)
         h_new = A_bar * h_prev + Bx
-
-        y = (h_new * Cmat.unsqueeze(1)).sum(-1) + self.D * x_in  # (B, inner_dim)
+        y = (h_new * Cmat.unsqueeze(1)).sum(-1) + self.D * x_in
         y = y * F.silu(gate)
-        return self.out_proj(y), h_new
+        return (self.out_proj(y), h_new)
 
 
 class MambaBlock(nn.Module):
+
     def __init__(self, dim: int, state_dim: int = 16, expand: int = 2):
         super().__init__()
         self.norm = RMSNorm(dim)
@@ -158,12 +102,15 @@ class MambaBlock(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x + self.ssm(self.norm(x))
 
-    def step(self, x_t: torch.Tensor, h_prev: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]:
+    def step(
+        self, x_t: torch.Tensor, h_prev: torch.Tensor | None
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         y, h_new = self.ssm.step(self.norm(x_t), h_prev)
-        return x_t + y, h_new
+        return (x_t + y, h_new)
 
 
 class ChessMamba(nn.Module):
+
     def __init__(
         self,
         dim: int = 256,
@@ -176,32 +123,25 @@ class ChessMamba(nn.Module):
         super().__init__()
         self.dim = dim
         self.max_plies = max_plies
-        # gradient checkpointing recomputes each block's forward during backward instead of
-        # retaining all depth * seq_len intermediate scan tensors at once -- without this,
-        # backprop through a ~96-step sequential scan stacked depth=10 deep OOMs even an 80GB
-        # GPU (verified: batch 1024 at dim=384/depth=10 exhausted 80GB with checkpoint off).
         self.use_checkpoint = use_checkpoint
-
         self.from_embed = nn.Embedding(64, dim)
         self.to_embed = nn.Embedding(64, dim)
         self.promo_embed = nn.Embedding(NUM_PROMO, dim)
-        self.pos_embed = nn.Embedding(max_plies + 1, dim)  # +1 for the <start> position
+        self.pos_embed = nn.Embedding(max_plies + 1, dim)
         self.start_token = nn.Parameter(torch.zeros(1, 1, dim))
-
         self.blocks = nn.ModuleList([MambaBlock(dim, state_dim, expand) for _ in range(depth)])
         self.norm_f = RMSNorm(dim)
-
         self.policy_head = nn.Linear(dim, NUM_FROM_TO)
         self.promo_head = nn.Linear(dim, NUM_PROMO)
         self.value_head = nn.Linear(dim, 1)
-
         nn.init.normal_(self.from_embed.weight, std=0.02)
         nn.init.normal_(self.to_embed.weight, std=0.02)
         nn.init.normal_(self.promo_embed.weight, std=0.02)
         nn.init.normal_(self.pos_embed.weight, std=0.02)
 
-    def embed_moves(self, from_ids: torch.Tensor, to_ids: torch.Tensor, promo_ids: torch.Tensor) -> torch.Tensor:
-        """from_ids/to_ids/promo_ids: (B, L) move history, L can be 0 (empty game)."""
+    def embed_moves(
+        self, from_ids: torch.Tensor, to_ids: torch.Tensor, promo_ids: torch.Tensor
+    ) -> torch.Tensor:
         B_ = from_ids.shape[0]
         start = self.start_token.expand(B_, 1, -1)
         if from_ids.shape[1] == 0:
@@ -212,17 +152,13 @@ class ChessMamba(nn.Module):
         positions = torch.arange(tok.shape[1], device=tok.device).unsqueeze(0)
         return tok + self.pos_embed(positions)
 
-    def forward(self, from_ids: torch.Tensor, to_ids: torch.Tensor, promo_ids: torch.Tensor, lengths: torch.Tensor | None = None):
-        """Returns (policy_logits (B,4096), promo_logits (B,5), value (B,1)).
-
-        `lengths` (B,), if given, is the real (unpadded) move-history length per example --
-        needed so a batch can mix examples of different history lengths: sequences are
-        right-padded to a common L_max, and since the scan is causal, padding after an
-        example's real length never affects the output gathered at that example's own
-        position (index == lengths[i], since index 0 is the <start> token). No attention
-        mask needed, unlike a transformer. If None, the last position is used (for live
-        incremental play where the input is exactly the current position, no padding).
-        """
+    def forward(
+        self,
+        from_ids: torch.Tensor,
+        to_ids: torch.Tensor,
+        promo_ids: torch.Tensor,
+        lengths: torch.Tensor | None = None,
+    ):
         x = self.embed_moves(from_ids, to_ids, promo_ids)
         for block in self.blocks:
             if self.use_checkpoint and self.training:
@@ -231,21 +167,14 @@ class ChessMamba(nn.Module):
                 x = block(x)
         x = self.norm_f(x)
         if lengths is None:
-            pooled = x[:, -1]  # (B, dim)
+            pooled = x[:, -1]
         else:
             idx = lengths.view(-1, 1, 1).expand(-1, 1, x.shape[-1])
             pooled = x.gather(1, idx).squeeze(1)
         policy_logits = self.policy_head(pooled)
         promo_logits = self.promo_head(pooled)
         value = torch.tanh(self.value_head(pooled))
-        return policy_logits, promo_logits, value
-
-    # -- incremental (single-move) inference path, for live/search play --
-    # Search visits many positions per move (up to top_k^depth); recomputing the full
-    # windowed forward pass at every node was the actual search-speed bottleneck. These
-    # methods carry per-block hidden state across moves so each new node costs O(1)
-    # (one token through `depth` layers) instead of O(window length). State is
-    # (next_position_index, [per-block hidden state or None]).
+        return (policy_logits, promo_logits, value)
 
     @torch.no_grad()
     def init_incremental(self, device: torch.device | str = "cpu"):
@@ -257,14 +186,14 @@ class ChessMamba(nn.Module):
             block_states.append(s)
         x = self.norm_f(x)
         outputs = self._heads(x)
-        return (1, block_states), outputs
+        return ((1, block_states), outputs)
 
     @torch.no_grad()
     def step_move(self, from_sq: int, to_sq: int, promo_id: int, state):
         pos_idx, block_states = state
         device = self.from_embed.weight.device
         pos_idx_clamped = min(pos_idx, self.max_plies)
-        idx = lambda v: torch.tensor([v], device=device)  # noqa: E731
+        idx = lambda v: torch.tensor([v], device=device)
         x = (
             self.from_embed(idx(from_sq))
             + self.to_embed(idx(to_sq))
@@ -277,33 +206,36 @@ class ChessMamba(nn.Module):
             new_block_states.append(ns)
         x = self.norm_f(x)
         outputs = self._heads(x)
-        return (pos_idx + 1, new_block_states), outputs
+        return ((pos_idx + 1, new_block_states), outputs)
 
     @torch.no_grad()
-    def build_incremental_state(self, from_list: list[int], to_list: list[int], promo_list: list[int], device: torch.device | str = "cpu"):
-        """Rebuild state from a (already-windowed, len <= max_plies) move list -- exactly
-        matching what a windowed `forward()` call would see. Used once per real move
-        played (cheap: len(list) O(1) steps), not per search node."""
+    def build_incremental_state(
+        self,
+        from_list: list[int],
+        to_list: list[int],
+        promo_list: list[int],
+        device: torch.device | str = "cpu",
+    ):
         state, outputs = self.init_incremental(device)
         for f, t, p in zip(from_list, to_list, promo_list):
             state, outputs = self.step_move(f, t, p, state)
-        return state, outputs
+        return (state, outputs)
 
     def _heads(self, x: torch.Tensor):
         policy_logits = self.policy_head(x)
         promo_logits = self.promo_head(x)
         value = torch.tanh(self.value_head(x))
-        return policy_logits, promo_logits, value
+        return (policy_logits, promo_logits, value)
 
 
 def count_params(model: nn.Module) -> int:
-    return sum(p.numel() for p in model.parameters())
+    return sum((p.numel() for p in model.parameters()))
 
 
 if __name__ == "__main__":
     m = ChessMamba(dim=256, depth=8)
     print(f"params: {count_params(m):,}")
-    B_, L = 4, 20
+    B_, L = (4, 20)
     from_ids = torch.randint(0, 64, (B_, L))
     to_ids = torch.randint(0, 64, (B_, L))
     promo_ids = torch.zeros(B_, L, dtype=torch.long)

search.py

@@ -1,70 +1,28 @@
-"""Shallow policy-guided negamax search on top of ChessMamba.
-
-Why search at all, given a trained policy/value net: a small, data-limited model
-playing pure policy-argmax would blunder tactically (DeepMind's search-free
-transformer needed 270M params + 10M Stockfish-annotated positions to get away
-with zero search -- we have neither). Plan: expand the top-K policy-ranked moves
-at each node, shallow alpha-beta (small fixed depth) with the value head as leaf
-eval. This protects the Elo/win-loss score, one of the tournament's three scoring
-axes alongside uniqueness.
-
-Search visits many nodes per move (up to top_k^depth). Each node used to pay for a
-full windowed forward pass (recomputing the whole move-history scan from scratch,
-`model.py`'s pre-incremental design) -- the actual search-speed bottleneck. Nodes
-here instead carry `model.py`'s incremental (state, outputs) pair: computing a
-child node is one O(1) step through the model, not an O(window length) rescan.
-Falls back to a random legal move if no checkpoint is present, so the engine is
-always submittable regardless of how training goes under the deadline.
-"""
-
 from __future__ import annotations
-
 import os
 import random
 import time
-
 import torch
 
-# search does single-position (batch=1) sequential incremental steps -- torch's default
-# intra-op multi-threading mainly benefits large batched tensor ops, not this workload
-# (confirmed by the CPU-beats-CUDA benchmark: latency-bound, not throughput-bound), so
-# there's little to lose here. There's a lot to lose from NOT setting this though: two
-# engine processes on the same test box (e.g. testing our own engine against itself, or
-# any match where both sides run locally) each default to using every core, causing the
-# same thread-oversubscription/contention bug already fixed in selfplay_finetune.py --
-# caught when a diagnostic match silently ran ~15x slower than its movetime budget implied.
 torch.set_num_threads(1)
-
 import chess
-
 from chess_io import move_to_ids, PROMO_INDEX
 from model import ChessMamba, MAX_PLIES
 
 DEFAULT_TOP_K = 10
-DEFAULT_MAX_DEPTH = 6  # raised from 3 now that each node is O(1) instead of O(window length)
+DEFAULT_MAX_DEPTH = 6
 DEFAULT_MOVETIME_S = 3.0
-DEFAULT_MAX_QDEPTH = 6  # quiescence extension depth beyond the main search horizon
-DEFAULT_QS_TOP_K = 6  # capture/promotion candidates examined per quiescence node
-
-# fallback budgeting when the GUI sends a real clock (wtime/btime/winc/binc) instead of
-# a fixed movetime -- simple "assume ~N moves left" heuristic, same idea most small
-# engines use. Only engages when movetime is absent; a fixed movetime always wins.
+DEFAULT_MAX_QDEPTH = 6
+DEFAULT_QS_TOP_K = 6
 CLOCK_MOVES_DIVISOR = 30
 MIN_CLOCK_MOVETIME_S = 0.05
-# hard ceiling applied AFTER any endgame extension, regardless of source (movetime,
-# clock heuristic, or default) -- Ender's stated tournament rule is "no more than 20
-# seconds" per move; without this, a movetime near 20000ms plus the 1.6x endgame
-# extension would silently blow past that (e.g. 20s * 1.6 = 32s).
 HARD_MOVETIME_CEILING_S = 19.0
-
-# rough material count, used ONLY to decide how much search time an endgame position
-# deserves -- a resource-allocation heuristic, not a position evaluation (the value head
-# still does all actual evaluating). Standard piece-value scale, nothing unusual.
 _MATERIAL_VALUE = {chess.PAWN: 1, chess.KNIGHT: 3, chess.BISHOP: 3, chess.ROOK: 5, chess.QUEEN: 9}
-_ENDGAME_MATERIAL_THRESHOLD = 14  # combined non-king material below this = "endgame"
+_ENDGAME_MATERIAL_THRESHOLD = 14
 
 
 class Searcher:
+
     def __init__(
         self,
         checkpoint_path: str | None = None,
@@ -75,17 +33,7 @@ class Searcher:
         device: str = "cpu",
         policy_only: bool = False,
     ):
-        # policy_only: skip negamax entirely, just play the top policy-ranked legal move.
-        # Used to diagnose whether the search is actually a genuine improvement over raw
-        # policy -- expert iteration (self-play fine-tuning) only helps if it is; if search
-        # doesn't reliably beat policy-only, self-play is just distilling the network's own
-        # blind spots back into itself, not a real correction.
         self.policy_only = policy_only
-        # NOT auto-detecting/defaulting to cuda even when available: search does many
-        # single-position (batch=1) incremental steps, and un-benchmarked GPU use for
-        # tiny sequential calls like this is as likely to be slower (kernel-launch and
-        # host<->device transfer overhead per node) as faster. Only flip the default
-        # after actually measuring it on real target-class hardware -- see PLAN.md.
         self.top_k = top_k
         self.max_depth = max_depth
         self.max_qdepth = max_qdepth
@@ -99,7 +47,7 @@ class Searcher:
             self.model.load_state_dict(ckpt["model"])
             self.model.eval()
             self.model.to(self.device)
-        self.root_node = None  # (state, outputs), see model.py's incremental interface
+        self.root_node = None
         self._root_key: tuple | None = None
 
     def reset(self) -> None:
@@ -107,24 +55,21 @@ class Searcher:
         self._root_key = None
 
     def sync(self, board: chess.Board) -> None:
-        """Rebuild the root node from the board's actual move history, matching what a
-        windowed forward() call would see. Cheap (<=MAX_PLIES incremental steps), and
-        cached against the move list so redundant calls for the same position are free."""
         if self.model is None:
             return
         window = board.move_stack[-MAX_PLIES:]
-        key = tuple(m.uci() for m in window)
+        key = tuple((m.uci() for m in window))
         if key == self._root_key:
             return
         self._root_key = key
         from_list = [m.from_square for m in window]
         to_list = [m.to_square for m in window]
         promo_list = [move_to_ids(m)[1] for m in window]
-        self.root_node = self.model.build_incremental_state(from_list, to_list, promo_list, device=self.device)
+        self.root_node = self.model.build_incremental_state(
+            from_list, to_list, promo_list, device=self.device
+        )
 
     def _expand(self, board: chess.Board, node, move: chess.Move):
-        """Pushes `move` on `board` and returns the child node. Caller must board.pop()
-        when done (the state itself needs no explicit undo -- it's just discarded)."""
         state, _ = node
         board.push(move)
         promo_id = PROMO_INDEX[move.promotion]
@@ -134,30 +79,28 @@ class Searcher:
     @staticmethod
     def _is_endgame(board: chess.Board) -> bool:
         total = sum(
-            _MATERIAL_VALUE.get(p.piece_type, 0)
-            for p in board.piece_map().values()
-            if p.piece_type != chess.KING
+            (
+                _MATERIAL_VALUE.get(p.piece_type, 0)
+                for p in board.piece_map().values()
+                if p.piece_type != chess.KING
+            )
         )
         return total < _ENDGAME_MATERIAL_THRESHOLD
 
     def _policy_value(self, board: chess.Board, node):
-        """Return (priors: dict[chess.Move -> float], value: float from side-to-move perspective)."""
         if board.is_checkmate():
-            return {}, -1.0
+            return ({}, -1.0)
         if board.is_stalemate() or board.is_insufficient_material() or board.can_claim_draw():
-            return {}, 0.0
-
+            return ({}, 0.0)
         legal = list(board.legal_moves)
         if not legal:
-            return {}, 0.0
+            return ({}, 0.0)
         if self.model is None:
             uniform = 1.0 / len(legal)
-            return {m: uniform for m in legal}, 0.0
-
+            return ({m: uniform for m in legal}, 0.0)
         _, (policy_logits, promo_logits, value) = node
         policy_logits = policy_logits[0]
         promo_logits = promo_logits[0]
-
         scores = {}
         for mv in legal:
             from_to_id = mv.from_square * 64 + mv.to_square
@@ -165,24 +108,19 @@ class Searcher:
             if mv.promotion is not None:
                 score += promo_logits[PROMO_INDEX[mv.promotion]].item()
             scores[mv] = score
-
         max_score = max(scores.values())
         exp_scores = {m: torch.exp(torch.tensor(s - max_score)).item() for m, s in scores.items()}
         total = sum(exp_scores.values()) or 1.0
         priors = {m: s / total for m, s in exp_scores.items()}
-        return priors, value.item()
+        return (priors, value.item())
 
-    def _quiescence(self, board: chess.Board, node, alpha: float, beta: float, deadline: float, qdepth: int = 0) -> float:
-        """Extends search through captures/promotions only, past the main search's fixed
-        horizon, so a leaf doesn't get statically evaluated mid-capture-sequence (the
-        classic horizon-effect blunder for any fixed-depth search). Move ordering reuses
-        the same policy priors the rest of the search uses -- no separate hand-crafted
-        piece-value ordering, consistent with evaluating everything through the net."""
+    def _quiescence(
+        self, board: chess.Board, node, alpha: float, beta: float, deadline: float, qdepth: int = 0
+    ) -> float:
         if board.is_checkmate():
             return -1.0
         if board.is_stalemate() or board.is_insufficient_material() or board.can_claim_draw():
             return 0.0
-
         priors, stand_pat = self._policy_value(board, node)
         if time.monotonic() > deadline or qdepth >= self.max_qdepth:
             return stand_pat
@@ -190,12 +128,10 @@ class Searcher:
             return beta
         if stand_pat > alpha:
             alpha = stand_pat
-
         loud = [mv for mv in priors if board.is_capture(mv) or mv.promotion is not None]
         if not loud:
             return alpha
         loud.sort(key=lambda mv: -priors[mv])
-
         for mv in loud[: self.qs_top_k]:
             child = self._expand(board, node, mv)
             score = -self._quiescence(board, child, -beta, -alpha, deadline, qdepth + 1)
@@ -208,7 +144,9 @@ class Searcher:
                 break
         return alpha
 
-    def _negamax(self, board: chess.Board, node, depth: int, alpha: float, beta: float, deadline: float) -> float:
+    def _negamax(
+        self, board: chess.Board, node, depth: int, alpha: float, beta: float, deadline: float
+    ) -> float:
         if board.is_checkmate():
             return -1.0
         if board.is_stalemate() or board.is_insufficient_material() or board.can_claim_draw():
@@ -218,12 +156,10 @@ class Searcher:
             return value
         if depth == 0:
             return self._quiescence(board, node, alpha, beta, deadline)
-
         priors, _ = self._policy_value(board, node)
         if not priors:
             return 0.0
         ranked = sorted(priors.items(), key=lambda kv: -kv[1])[: self.top_k]
-
         best = -float("inf")
         for mv, _ in ranked:
             child = self._expand(board, node, mv)
@@ -241,39 +177,37 @@ class Searcher:
         legal = list(board.legal_moves)
         if not legal:
             return None
-
         movetime_s = DEFAULT_MOVETIME_S
         if "movetime" in go_tokens:
             movetime_s = int(go_tokens[go_tokens.index("movetime") + 1]) / 1000.0
         elif "wtime" in go_tokens or "btime" in go_tokens:
             time_key = "wtime" if board.turn == chess.WHITE else "btime"
             inc_key = "winc" if board.turn == chess.WHITE else "binc"
-            own_time_ms = int(go_tokens[go_tokens.index(time_key) + 1]) if time_key in go_tokens else 0
-            own_inc_ms = int(go_tokens[go_tokens.index(inc_key) + 1]) if inc_key in go_tokens else 0
+            own_time_ms = (
+                int(go_tokens[go_tokens.index(time_key) + 1]) if time_key in go_tokens else 0
+            )
+            own_inc_ms = (
+                int(go_tokens[go_tokens.index(inc_key) + 1]) if inc_key in go_tokens else 0
+            )
             movetime_s = own_time_ms / 1000.0 / CLOCK_MOVES_DIVISOR + own_inc_ms / 1000.0
             movetime_s = max(MIN_CLOCK_MOVETIME_S, movetime_s)
         if self._is_endgame(board):
-            movetime_s *= 1.6  # low-material positions are where a small model's static eval is weakest
+            movetime_s *= 1.6
         movetime_s = min(movetime_s, HARD_MOVETIME_CEILING_S)
         deadline = time.monotonic() + movetime_s
-
         if self.model is None:
             return random.choice(legal)
-
         self.sync(board)
         node = self.root_node
-
         priors, _ = self._policy_value(board, node)
         ranked = sorted(priors.items(), key=lambda kv: -kv[1])[: self.top_k]
-
         best_move = ranked[0][0] if ranked else random.choice(legal)
         if self.policy_only:
             return best_move
         best_score = -float("inf")
-
         depth = 1
         while depth <= self.max_depth and time.monotonic() < deadline:
-            alpha, beta = -float("inf"), float("inf")
+            alpha, beta = (-float("inf"), float("inf"))
             current_best_move = None
             current_best_score = -float("inf")
             for mv, _ in ranked:
@@ -288,7 +222,6 @@ class Searcher:
                 if time.monotonic() > deadline:
                     break
             if current_best_move is not None and time.monotonic() <= deadline:
-                best_move, best_score = current_best_move, current_best_score
+                best_move, best_score = (current_best_move, current_best_score)
             depth += 1
-
         return best_move

training/data_pipeline.py

@@ -1,45 +1,30 @@
-"""Streams Lichess monthly rated-game dumps straight off the network (no full
-download-then-decompress -- decompresses and parses on the fly, stops once enough
-games are collected) and packs them into fixed-shape numpy shards for training.
-
-Each training example = one ply: (move history up to MAX_PLIES back, the move
-actually played, the eventual game result from the mover's perspective). This
-mirrors exactly what search.py's encode_history() builds at inference time.
-
-Run on the box with real network access / disk, not necessarily locally:
-    python3 data_pipeline.py --out-dir /root/chess/data --target-games 500000
-"""
-
 from __future__ import annotations
-
 import argparse
 import io
 import os
 import time
 import urllib.request
-
 import chess
 import chess.pgn
 import numpy as np
 import zstandard
-
 from chess_io import move_to_ids
 from model import MAX_PLIES
 
 BASE_URL = "https://database.lichess.org/standard/lichess_db_standard_rated_{}.pgn.zst"
 MIN_ELO = 1800
 MIN_PLIES = 10
-SHARD_SIZE = 200_000
-SOFT_K = 10  # candidate slots for distributional self-play targets (see ShardWriter.add)
+SHARD_SIZE = 200000
+SOFT_K = 10
 
 
 def candidate_months(start_year: int = 2026, start_month: int = 7, count: int = 30):
-    y, m = start_year, start_month
+    y, m = (start_year, start_month)
     for _ in range(count):
         yield f"{y:04d}-{m:02d}"
         m -= 1
         if m == 0:
-            m, y = 12, y - 1
+            m, y = (12, y - 1)
 
 
 def find_available_month() -> str:
@@ -85,7 +70,6 @@ def game_is_ok(game: chess.pgn.Game) -> bool:
 
 
 def result_value(result: str, side_to_move: bool) -> int:
-    """side_to_move: chess.WHITE (True) or chess.BLACK (False). Returns +1/0/-1 from that side's view."""
     if result == "1/2-1/2":
         return 0
     white_won = result == "1-0"
@@ -99,10 +83,8 @@ def game_to_examples(game: chess.pgn.Game):
     hist_from: list[int] = []
     hist_to: list[int] = []
     hist_promo: list[int] = []
-
     if len(list(game.mainline_moves())) < MIN_PLIES:
         return
-
     for move in game.mainline_moves():
         side_to_move = board.turn
         from_to_id, promo_id = move_to_ids(move)
@@ -111,7 +93,6 @@ def game_to_examples(game: chess.pgn.Game):
         window_promo = hist_promo[-MAX_PLIES:]
         value = result_value(result, side_to_move)
         yield (window_from, window_to, window_promo, from_to_id, promo_id, value)
-
         hist_from.append(move.from_square)
         hist_to.append(move.to_square)
         hist_promo.append(promo_id)
@@ -119,6 +100,7 @@ def game_to_examples(game: chess.pgn.Game):
 
 
 class ShardWriter:
+
     def __init__(self, out_dir: str, shard_size: int = SHARD_SIZE):
         os.makedirs(out_dir, exist_ok=True)
         self.out_dir = out_dir
@@ -135,20 +117,21 @@ class ShardWriter:
         self.target_from_to = np.zeros((n,), dtype=np.int16)
         self.target_promo = np.zeros((n,), dtype=np.int8)
         self.values = np.zeros((n,), dtype=np.int8)
-        # distributional policy target (see class docstring on `add`) -- degenerate
-        # one-hot by default so ordinary hard-label writes (data_pipeline.py's own
-        # game_to_examples) are unaffected and unambiguous to a soft-aware loss.
         self.soft_move_ids = np.full((n, SOFT_K), -1, dtype=np.int16)
         self.soft_weights = np.zeros((n, SOFT_K), dtype=np.float32)
         self.count = 0
 
-    def add(self, window_from, window_to, window_promo, target_from_to, target_promo, value, soft_move_ids=None, soft_weights=None):
-        """soft_move_ids/soft_weights (optional): up to SOFT_K (from_to_id, weight) pairs
-        giving a distribution over candidate moves, e.g. a search's score distribution
-        over root candidates -- richer self-play training signal than a single hard
-        label. Omit for ordinary hard-label examples (data_pipeline.py's default path);
-        target_from_to is then stored as a degenerate one-hot distribution instead, so
-        pretrain.py's loss can treat every example uniformly."""
+    def add(
+        self,
+        window_from,
+        window_to,
+        window_promo,
+        target_from_to,
+        target_promo,
+        value,
+        soft_move_ids=None,
+        soft_weights=None,
+    ):
         i = self.count
         L = len(window_from)
         if L:
@@ -194,18 +177,17 @@ class ShardWriter:
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--out-dir", default="/root/chess/data")
-    ap.add_argument("--target-games", type=int, default=500_000)
-    ap.add_argument("--target-examples", type=int, default=0, help="0 = no cap, stop on target-games instead")
+    ap.add_argument("--target-games", type=int, default=500000)
+    ap.add_argument(
+        "--target-examples", type=int, default=0, help="0 = no cap, stop on target-games instead"
+    )
     args = ap.parse_args()
-
     url = find_available_month()
     writer = ShardWriter(args.out_dir)
-
     kept_games = 0
     seen_games = 0
     total_examples = 0
     t0 = time.monotonic()
-
     for game in stream_games(url):
         seen_games += 1
         if not game_is_ok(game):
@@ -214,19 +196,15 @@ def main():
             writer.add(*ex)
             total_examples += 1
         kept_games += 1
-
         if kept_games % 2000 == 0:
             dt = time.monotonic() - t0
             print(
-                f"[data_pipeline] kept={kept_games} seen={seen_games} examples={total_examples} "
-                f"elapsed={dt:.0f}s rate={kept_games/dt:.1f} games/s"
+                f"[data_pipeline] kept={kept_games} seen={seen_games} examples={total_examples} elapsed={dt:.0f}s rate={kept_games / dt:.1f} games/s"
             )
-
         if kept_games >= args.target_games:
             break
         if args.target_examples and total_examples >= args.target_examples:
             break
-
     writer.flush()
     print(f"[data_pipeline] DONE kept={kept_games} seen={seen_games} examples={total_examples}")
 

training/match.py

@@ -1,24 +1,13 @@
-"""Minimal UCI-vs-UCI match driver, for informal calibration -- not a submission
-artifact. Runs two engines as subprocesses speaking the UCI protocol over
-stdin/stdout, alternates moves on a shared board, reports the result. Useful for
-sanity-checking a checkpoint against a known reference (e.g. Sunfish) or against
-itself, since we have no access to a real rating pool before the tournament.
-
-Usage:
-    python3 match.py --white "python3 engine_uci.py" --black "python3 sunfish_uci.py" --movetime 200
-"""
-
 from __future__ import annotations
-
 import argparse
 import shlex
 import subprocess
 import time
-
 import chess
 
 
 class UciEngine:
+
     def __init__(self, command: str, cwd: str | None = None):
         self.proc = subprocess.Popen(
             shlex.split(command),
@@ -74,12 +63,13 @@ class UciEngine:
             self.proc.kill()
 
 
-def play_game(white_cmd: str, black_cmd: str, movetime_ms: int, max_plies: int = 300) -> tuple[str, list[str]]:
+def play_game(
+    white_cmd: str, black_cmd: str, movetime_ms: int, max_plies: int = 300
+) -> tuple[str, list[str]]:
     white = UciEngine(white_cmd)
     black = UciEngine(black_cmd)
     white.new_game()
     black.new_game()
-
     board = chess.Board()
     moves_uci: list[str] = []
     result = "*"
@@ -90,7 +80,9 @@ def play_game(white_cmd: str, black_cmd: str, movetime_ms: int, max_plies: int =
             move = chess.Move.from_uci(mv_uci)
             if move not in board.legal_moves:
                 result = "0-1" if board.turn == chess.WHITE else "1-0"
-                print(f"[match] ILLEGAL MOVE {mv_uci} by {'white' if board.turn else 'black'} -- forfeit")
+                print(
+                    f"[match] ILLEGAL MOVE {mv_uci} by {('white' if board.turn else 'black')} -- forfeit"
+                )
                 break
             board.push(move)
             moves_uci.append(mv_uci)
@@ -98,11 +90,11 @@ def play_game(white_cmd: str, black_cmd: str, movetime_ms: int, max_plies: int =
                 result = board.result()
                 break
         else:
-            result = "1/2-1/2"  # ply cap reached, call it a draw for calibration purposes
+            result = "1/2-1/2"
     finally:
         white.close()
         black.close()
-    return result, moves_uci
+    return (result, moves_uci)
 
 
 def main():
@@ -112,7 +104,6 @@ def main():
     ap.add_argument("--movetime", type=int, default=300)
     ap.add_argument("--games", type=int, default=1)
     args = ap.parse_args()
-
     tally = {"white": 0, "black": 0, "draw": 0}
     for g in range(args.games):
         w, b = (args.white, args.black) if g % 2 == 0 else (args.black, args.white)
@@ -124,7 +115,9 @@ def main():
             tally["black" if g % 2 == 0 else "white"] += 1
         else:
             tally["draw"] += 1
-    print(f"[summary] --white wins={tally['white']} --black wins={tally['black']} draws={tally['draw']}")
+    print(
+        f"[summary] --white wins={tally['white']} --black wins={tally['black']} draws={tally['draw']}"
+    )
 
 
 if __name__ == "__main__":

training/pretrain.py

@@ -1,39 +1,25 @@
-"""Supervised pretraining for ChessMamba: masked-legal next-move cross-entropy +
-game-outcome value regression, over shards produced by data_pipeline.py.
-
-Deliberately plain AdamW, not a custom optimizer -- this is a ~3.5-day deadline
-project, not a from-scratch-optimizer research effort like the LM projects.
-"""
-
 from __future__ import annotations
-
 import argparse
 import glob
 import os
 import time
-
 import numpy as np
 import torch
 import torch.nn.functional as F
-
 from data_pipeline import SOFT_K
 from model import ChessMamba, MAX_PLIES
 
 
 class ShardDataset:
-    """Loads all shards fully into RAM (data is small enough -- tens of GB at most
-    against 188GB host RAM) and hands out shuffled batches with per-batch length
-    truncation (batches are padded to their own max real length, not the global
-    MAX_PLIES, since the scan cost is O(L))."""
 
     def __init__(self, data_dir: str, val_fraction: float = 0.01):
-        # recursive glob: parallel self-play workers each write shard_*.npz into their own
-        # subdirectory (avoids filename collisions between workers writing concurrently)
         paths = sorted(glob.glob(os.path.join(data_dir, "**", "shard_*.npz"), recursive=True))
         if not paths:
             raise RuntimeError(f"no shards found in {data_dir}")
-        from_ids, to_ids, promo_ids, lengths, target_from_to, target_promo, values = ([] for _ in range(7))
-        soft_move_ids, soft_weights = [], []
+        from_ids, to_ids, promo_ids, lengths, target_from_to, target_promo, values = (
+            [] for _ in range(7)
+        )
+        soft_move_ids, soft_weights = ([], [])
         for p in paths:
             d = np.load(p)
             from_ids.append(d["from_ids"])
@@ -47,9 +33,6 @@ class ShardDataset:
                 soft_move_ids.append(d["soft_move_ids"])
                 soft_weights.append(d["soft_weights"])
             else:
-                # shards written before the soft-target field existed: synthesize the
-                # equivalent degenerate one-hot distribution so the loss can treat every
-                # example uniformly regardless of which stage produced it.
                 n_ex = len(d["target_from_to"])
                 one_hot_ids = np.full((n_ex, SOFT_K), -1, dtype=np.int16)
                 one_hot_ids[:, 0] = d["target_from_to"]
@@ -57,7 +40,6 @@ class ShardDataset:
                 one_hot_w[:, 0] = 1.0
                 soft_move_ids.append(one_hot_ids)
                 soft_weights.append(one_hot_w)
-
         self.from_ids = np.concatenate(from_ids)
         self.to_ids = np.concatenate(to_ids)
         self.promo_ids = np.concatenate(promo_ids)
@@ -67,20 +49,17 @@ class ShardDataset:
         self.values = np.concatenate(values)
         self.soft_move_ids = np.concatenate(soft_move_ids)
         self.soft_weights = np.concatenate(soft_weights)
-
         n = len(self.lengths)
         rng = np.random.default_rng(0)
         perm = rng.permutation(n)
         n_val = max(1, int(n * val_fraction))
         self.val_idx = perm[:n_val]
         self.train_idx = perm[n_val:]
-        print(f"[data] loaded {n:,} examples from {len(paths)} shards ({n - n_val:,} train / {n_val:,} val)")
+        print(
+            f"[data] loaded {n:,} examples from {len(paths)} shards ({n - n_val:,} train / {n_val:,} val)"
+        )
 
     def get_batch(self, idx: np.ndarray, device: torch.device, fixed_len: int | None = None):
-        # fixed_len pins every batch to the same shape (full MAX_PLIES by default) --
-        # required under torch.compile, which recompiles (expensive, ~60s) per distinct
-        # shape it sees; truncating to each batch's own max real length would otherwise
-        # trigger a recompile storm across the many different lengths in the data.
         max_len = fixed_len if fixed_len is not None else int(self.lengths[idx].max())
         from_ids = torch.from_numpy(self.from_ids[idx][:, :max_len].astype(np.int64)).to(device)
         to_ids = torch.from_numpy(self.to_ids[idx][:, :max_len].astype(np.int64)).to(device)
@@ -91,47 +70,57 @@ class ShardDataset:
         values = torch.from_numpy(self.values[idx].astype(np.float32)).to(device)
         soft_move_ids = torch.from_numpy(self.soft_move_ids[idx].astype(np.int64)).to(device)
         soft_weights = torch.from_numpy(self.soft_weights[idx].astype(np.float32)).to(device)
-        return from_ids, to_ids, promo_ids, lengths, target_from_to, target_promo, values, soft_move_ids, soft_weights
+        return (
+            from_ids,
+            to_ids,
+            promo_ids,
+            lengths,
+            target_from_to,
+            target_promo,
+            values,
+            soft_move_ids,
+            soft_weights,
+        )
 
 
 def concat_batches(batch_a, batch_b):
-    """Concatenate two get_batch() tuples along the batch dim -- both were built with the
-    same fixed_len, so every tensor pair already shares its non-batch dimensions."""
-    return tuple(torch.cat([a, b], dim=0) for a, b in zip(batch_a, batch_b))
+    return tuple((torch.cat([a, b], dim=0) for a, b in zip(batch_a, batch_b)))
 
 
 def compute_loss(model: ChessMamba, batch, device: torch.device, value_loss_weight: float = 0.5):
-    from_ids, to_ids, promo_ids, lengths, target_from_to, target_promo, values, soft_move_ids, soft_weights = batch
+    (
+        from_ids,
+        to_ids,
+        promo_ids,
+        lengths,
+        target_from_to,
+        target_promo,
+        values,
+        soft_move_ids,
+        soft_weights,
+    ) = batch
     policy_logits, promo_logits, value = model(from_ids, to_ids, promo_ids, lengths=lengths)
-
-    # soft cross-entropy over up to SOFT_K candidate moves per example, weighted by
-    # soft_weights -- mathematically identical to F.cross_entropy(policy_logits,
-    # target_from_to) when the target is a degenerate one-hot (ordinary hard-label
-    # data), and a genuine distributional loss when it isn't (self-play fine-tuning
-    # targets, e.g. softmax-with-temperature over a search's root candidate scores).
-    # Padding slots (id -1) are clamped to a valid index for the gather but contribute
-    # nothing since their weight is 0.
     log_probs = F.log_softmax(policy_logits, dim=-1)
-    gathered = log_probs.gather(1, soft_move_ids.clamp(min=0))  # (B, SOFT_K)
+    gathered = log_probs.gather(1, soft_move_ids.clamp(min=0))
     policy_loss = -(soft_weights * gathered).sum(-1).mean()
-
     promo_mask = target_promo > 0
     if promo_mask.any():
         promo_loss = F.cross_entropy(promo_logits[promo_mask], target_promo[promo_mask])
     else:
         promo_loss = torch.tensor(0.0, device=device)
-
     value_loss = F.mse_loss(value.squeeze(-1), values)
-
     loss = policy_loss + 0.25 * promo_loss + value_loss_weight * value_loss
     with torch.no_grad():
         acc = (policy_logits.argmax(-1) == target_from_to).float().mean()
-    return loss, {
-        "policy_loss": policy_loss.item(),
-        "promo_loss": promo_loss.item() if promo_loss.requires_grad else float(promo_loss),
-        "value_loss": value_loss.item(),
-        "move_acc": acc.item(),
-    }
+    return (
+        loss,
+        {
+            "policy_loss": policy_loss.item(),
+            "promo_loss": promo_loss.item() if promo_loss.requires_grad else float(promo_loss),
+            "value_loss": value_loss.item(),
+            "move_acc": acc.item(),
+        },
+    )
 
 
 def main():
@@ -142,56 +131,75 @@ def main():
     ap.add_argument("--depth", type=int, default=10)
     ap.add_argument("--state-dim", type=int, default=16)
     ap.add_argument("--batch-size", type=int, default=1024)
-    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--lr", type=float, default=0.0003)
     ap.add_argument("--weight-decay", type=float, default=0.01)
-    ap.add_argument("--value-loss-weight", type=float, default=0.5, help="raised for phase 4: diagnosed bottleneck is the value head's material judgment specifically")
-    ap.add_argument("--steps", type=int, default=20_000)
+    ap.add_argument(
+        "--value-loss-weight",
+        type=float,
+        default=0.5,
+        help="raised for phase 4: diagnosed bottleneck is the value head's material judgment specifically",
+    )
+    ap.add_argument("--steps", type=int, default=20000)
     ap.add_argument("--warmup-steps", type=int, default=500)
     ap.add_argument("--log-every", type=int, default=50)
     ap.add_argument("--eval-every", type=int, default=500)
     ap.add_argument("--ckpt-every", type=int, default=1000)
     ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
-    ap.add_argument("--resume", default="", help="resume an interrupted run of THIS SAME phase (loads optimizer + step count too)")
-    ap.add_argument("--init-from", default="", help="warm-start a NEW phase from another checkpoint's weights only (fresh optimizer/step/LR schedule)")
-    ap.add_argument("--no-checkpoint", action="store_true", help="disable gradient checkpointing (faster per-step if it fits in memory)")
+    ap.add_argument(
+        "--resume",
+        default="",
+        help="resume an interrupted run of THIS SAME phase (loads optimizer + step count too)",
+    )
+    ap.add_argument(
+        "--init-from",
+        default="",
+        help="warm-start a NEW phase from another checkpoint's weights only (fresh optimizer/step/LR schedule)",
+    )
+    ap.add_argument(
+        "--no-checkpoint",
+        action="store_true",
+        help="disable gradient checkpointing (faster per-step if it fits in memory)",
+    )
     ap.add_argument("--no-compile", action="store_true", help="disable torch.compile")
-    ap.add_argument("--replay-data-dir", default="", help="mix in examples from this dir (e.g. the original supervised data) during fine-tuning, to avoid catastrophic forgetting of the base distribution")
-    ap.add_argument("--replay-fraction", type=float, default=0.4, help="fraction of each batch drawn from --replay-data-dir when set")
+    ap.add_argument(
+        "--replay-data-dir",
+        default="",
+        help="mix in examples from this dir (e.g. the original supervised data) during fine-tuning, to avoid catastrophic forgetting of the base distribution",
+    )
+    ap.add_argument(
+        "--replay-fraction",
+        type=float,
+        default=0.4,
+        help="fraction of each batch drawn from --replay-data-dir when set",
+    )
     args = ap.parse_args()
-
     os.makedirs(args.ckpt_dir, exist_ok=True)
     device = torch.device(args.device)
     if device.type == "cuda":
         torch.set_float32_matmul_precision("high")
-
     ds = ShardDataset(args.data_dir)
     replay_ds = ShardDataset(args.replay_data_dir) if args.replay_data_dir else None
-
     model_cfg = dict(dim=args.dim, depth=args.depth, state_dim=args.state_dim)
     model = ChessMamba(**model_cfg, use_checkpoint=not args.no_checkpoint).to(device)
-    # the parallel scan is many small elementwise/cat ops -- torch.compile fuses them into far
-    # fewer kernel launches. Measured impact: 7.8s/step -> 2.1s/step at dim=384/depth=10/batch=512.
     train_model = model if args.no_compile else torch.compile(model)
-    n_params = sum(p.numel() for p in model.parameters())
+    n_params = sum((p.numel() for p in model.parameters()))
     print(f"[model] {n_params:,} params, cfg={model_cfg}")
-
-    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95))
-
+    opt = torch.optim.AdamW(
+        model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95)
+    )
     start_step = 0
     if args.resume and os.path.exists(args.resume):
-        # true resume: same phase, same step budget/LR schedule, picking up an interrupted run
         ckpt = torch.load(args.resume, map_location=device, weights_only=False)
         model.load_state_dict(ckpt["model"])
         opt.load_state_dict(ckpt["optimizer"])
         start_step = ckpt.get("step", 0)
         print(f"[resume] loaded {args.resume} at step {start_step}")
     elif args.init_from and os.path.exists(args.init_from):
-        # warm-start: new phase (new dataset/step budget/LR schedule), only the weights carry
-        # over -- NOT optimizer state or step count, which would otherwise make `range(start_step,
-        # args.steps)` empty if this phase's step budget is smaller than the previous phase's.
         ckpt = torch.load(args.init_from, map_location=device, weights_only=False)
         model.load_state_dict(ckpt["model"])
-        print(f"[init-from] loaded weights from {args.init_from} (step {ckpt.get('step', '?')} of its own run), starting fresh optimizer/schedule at step 0")
+        print(
+            f"[init-from] loaded weights from {args.init_from} (step {ckpt.get('step', '?')} of its own run), starting fresh optimizer/schedule at step 0"
+        )
 
     def lr_at(step):
         if step < args.warmup_steps:
@@ -201,62 +209,67 @@ def main():
 
     rng = np.random.default_rng(1234)
     t0 = time.monotonic()
-
     replay_size = int(args.batch_size * args.replay_fraction) if replay_ds is not None else 0
     primary_size = args.batch_size - replay_size
     if replay_ds is not None:
-        print(f"[replay] mixing {replay_size}/{args.batch_size} examples per batch from {args.replay_data_dir}")
-
+        print(
+            f"[replay] mixing {replay_size}/{args.batch_size} examples per batch from {args.replay_data_dir}"
+        )
     for step in range(start_step, args.steps):
         for g in opt.param_groups:
             g["lr"] = lr_at(step)
-
         idx = rng.choice(ds.train_idx, size=primary_size, replace=False)
         batch = ds.get_batch(idx, device, fixed_len=MAX_PLIES)
         if replay_ds is not None:
             replay_idx = rng.choice(replay_ds.train_idx, size=replay_size, replace=False)
             replay_batch = replay_ds.get_batch(replay_idx, device, fixed_len=MAX_PLIES)
             batch = concat_batches(batch, replay_batch)
-
         model.train()
-        loss, stats = compute_loss(train_model, batch, device, value_loss_weight=args.value_loss_weight)
+        loss, stats = compute_loss(
+            train_model, batch, device, value_loss_weight=args.value_loss_weight
+        )
         opt.zero_grad(set_to_none=True)
         loss.backward()
         torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
         opt.step()
-
         if step % args.log_every == 0:
             dt = time.monotonic() - t0
             print(
-                f"step {step} loss {loss.item():.4f} policy {stats['policy_loss']:.4f} "
-                f"promo {stats['promo_loss']:.4f} value {stats['value_loss']:.4f} "
-                f"acc {stats['move_acc']:.3f} lr {lr_at(step):.2e} elapsed {dt:.0f}s"
+                f"step {step} loss {loss.item():.4f} policy {stats['policy_loss']:.4f} promo {stats['promo_loss']:.4f} value {stats['value_loss']:.4f} acc {stats['move_acc']:.3f} lr {lr_at(step):.2e} elapsed {dt:.0f}s"
             )
-
         if step % args.eval_every == 0 and step > 0:
             model.eval()
             with torch.no_grad():
-                val_idx = rng.choice(ds.val_idx, size=min(args.batch_size, len(ds.val_idx)), replace=False)
+                val_idx = rng.choice(
+                    ds.val_idx, size=min(args.batch_size, len(ds.val_idx)), replace=False
+                )
                 val_batch = ds.get_batch(val_idx, device, fixed_len=MAX_PLIES)
-                _, val_stats = compute_loss(train_model, val_batch, device, value_loss_weight=args.value_loss_weight)
-            print(f"  [val] step {step} policy {val_stats['policy_loss']:.4f} acc {val_stats['move_acc']:.3f}")
-
+                _, val_stats = compute_loss(
+                    train_model, val_batch, device, value_loss_weight=args.value_loss_weight
+                )
+            print(
+                f"  [val] step {step} policy {val_stats['policy_loss']:.4f} acc {val_stats['move_acc']:.3f}"
+            )
         if step % args.ckpt_every == 0 and step > 0:
             save_checkpoint(model, opt, step, model_cfg, args.ckpt_dir, "model.pt")
             print(f"  [ckpt] saved step {step}")
-
     save_checkpoint(model, opt, args.steps, model_cfg, args.ckpt_dir, "model.pt")
     print(f"[done] saved final checkpoint at step {args.steps}")
 
 
 def save_checkpoint(model, opt, step, model_cfg, ckpt_dir, name):
-    """Write-then-rename so a concurrent reader (e.g. testing the engine mid-training)
-    never sees a partially-written file -- os.replace is atomic on the same filesystem."""
     path = os.path.join(ckpt_dir, name)
     tmp_path = path + ".tmp"
-    torch.save({"model": model.state_dict(), "optimizer": opt.state_dict(), "step": step, "config": model_cfg}, tmp_path)
+    torch.save(
+        {
+            "model": model.state_dict(),
+            "optimizer": opt.state_dict(),
+            "step": step,
+            "config": model_cfg,
+        },
+        tmp_path,
+    )
     os.replace(tmp_path, path)
-    # keep a numbered snapshot too, so earlier checkpoints stay testable/comparable
     step_path = os.path.join(ckpt_dir, f"model_step{step}.pt")
     if not os.path.exists(step_path):
         import shutil

training/run_selfplay_parallel.py

@@ -1,15 +1,4 @@
-"""Fan `selfplay_finetune.py` out across multiple CPU processes -- self-play games are
-independent, and single-threaded generation (~1 game at a time) doesn't come close to
-using a 32-core box. Each worker writes to its own subdirectory (pretrain.py's
-ShardDataset globs recursively, so no merging step is needed afterward).
-
-Usage:
-    python3 run_selfplay_parallel.py --checkpoint /root/chess/ckpt/model.pt \
-        --out-dir /root/chess/selfplay_data --games 4000 --workers 28
-"""
-
 from __future__ import annotations
-
 import argparse
 import os
 import subprocess
@@ -25,7 +14,6 @@ def main():
     ap.add_argument("--movetime", type=int, default=200)
     ap.add_argument("--temperature-plies", type=int, default=10)
     args = ap.parse_args()
-
     per_worker = max(1, args.games // args.workers)
     procs = []
     logs = []
@@ -35,18 +23,28 @@ def main():
         os.makedirs(args.out_dir, exist_ok=True)
         log_f = open(log_path, "w")
         cmd = [
-            sys.executable, "selfplay_finetune.py",
-            "--checkpoint", args.checkpoint,
-            "--out-dir", worker_out,
-            "--games", str(per_worker),
-            "--movetime", str(args.movetime),
-            "--temperature-plies", str(args.temperature_plies),
+            sys.executable,
+            "selfplay_finetune.py",
+            "--checkpoint",
+            args.checkpoint,
+            "--out-dir",
+            worker_out,
+            "--games",
+            str(per_worker),
+            "--movetime",
+            str(args.movetime),
+            "--temperature-plies",
+            str(args.temperature_plies),
         ]
-        p = subprocess.Popen(cmd, stdout=log_f, stderr=subprocess.STDOUT, cwd=os.path.dirname(os.path.abspath(__file__)) or ".")
+        p = subprocess.Popen(
+            cmd,
+            stdout=log_f,
+            stderr=subprocess.STDOUT,
+            cwd=os.path.dirname(os.path.abspath(__file__)) or ".",
+        )
         procs.append(p)
         logs.append(log_f)
         print(f"[launcher] started worker {i} (pid {p.pid}), {per_worker} games -> {worker_out}")
-
     exit_codes = []
     for i, p in enumerate(procs):
         code = p.wait()
@@ -54,11 +52,14 @@ def main():
         print(f"[launcher] worker {i} exited with code {code}")
     for f in logs:
         f.close()
-
     failed = [i for i, c in enumerate(exit_codes) if c != 0]
     if failed:
-        print(f"[launcher] WARNING: workers failed: {failed} -- check worker*.log in {args.out_dir}")
-    print(f"[launcher] DONE: {args.workers} workers, {per_worker * args.workers} games total requested")
+        print(
+            f"[launcher] WARNING: workers failed: {failed} -- check worker*.log in {args.out_dir}"
+        )
+    print(
+        f"[launcher] DONE: {args.workers} workers, {per_worker * args.workers} games total requested"
+    )
 
 
 if __name__ == "__main__":

training/selfplay_finetune.py

@@ -1,116 +1,95 @@
-"""Generates self-play games with the current checkpoint's shallow search and packs
-them into shards in the exact same format data_pipeline.py produces -- so fine-tuning
-on them is just another pretrain.py run with --data-dir pointed here and --resume
-pointed at the pretrained checkpoint. No separate training loop needed.
-
-Policy target is a softmax-with-temperature distribution over the search's scored
-root candidates (not just the single chosen move) -- the AlphaZero-style motivation:
-a distributional target is a richer training signal than a hard label, which matters
-more here than it would with millions of games, since the self-play corpus is modest.
-Real AlphaZero uses MCTS visit counts specifically; ours uses negamax score-softmax
-over the same top-K candidates the rest of the search already ranks, which is the
-natural equivalent for a negamax (not MCTS) search -- not visit counts, but still a
-genuine distribution over multiple candidates instead of a single argmax label.
-"""
-
 from __future__ import annotations
-
 import argparse
 import math
 import random
 import time
-
 import torch
 
-# self-play runs many of these processes in parallel across cores (run_selfplay_parallel.py);
-# without this, each process's PyTorch defaults to using every core for its own tiny CPU
-# matmuls, so N workers oversubscribe to N*cores threads instead of using one core each.
 torch.set_num_threads(1)
-
 import chess
-
 from chess_io import move_to_ids
 from data_pipeline import ShardWriter
 from model import MAX_PLIES
 from search import Searcher
 
-
-SOFT_TARGET_TEMPERATURE = 0.5  # lower = closer to a hard argmax label, higher = flatter/softer
+SOFT_TARGET_TEMPERATURE = 0.5
 
 
 def scores_to_distribution(scored: list[tuple]) -> tuple[list[int], list[float]]:
-    """scored: [(move, negamax_score), ...] -> (from_to_ids, softmax weights)."""
     values = [s for _, s in scored]
     max_v = max(values)
     exp_vals = [math.exp((v - max_v) / SOFT_TARGET_TEMPERATURE) for v in values]
     total = sum(exp_vals) or 1.0
     ids = [mv.from_square * 64 + mv.to_square for mv, _ in scored]
     weights = [e / total for e in exp_vals]
-    return ids, weights
+    return (ids, weights)
 
 
-def play_selfplay_game(searcher: Searcher, movetime_ms: int, temperature_plies: int, max_plies: int = 200):
+def play_selfplay_game(
+    searcher: Searcher, movetime_ms: int, temperature_plies: int, max_plies: int = 200
+):
     board = chess.Board()
     searcher.reset()
-    searcher.sync(board)  # empty board -> root node = initial state
+    searcher.sync(board)
     node = searcher.root_node
     hist_from: list[int] = []
     hist_to: list[int] = []
     hist_promo: list[int] = []
-    ply_records = []  # (window_from, window_to, window_promo, target_from_to, target_promo, side_to_move, soft_ids, soft_weights)
-
+    ply_records = []
     for ply in range(max_plies):
         if board.is_game_over():
             break
-
         priors, _ = searcher._policy_value(board, node)
         if not priors:
             break
         ranked = sorted(priors.items(), key=lambda kv: -kv[1])[: searcher.top_k]
-
-        # score each candidate at the root via one ply of negamax, same mechanism choose_move uses.
-        # One shared deadline for the whole ply, not one per candidate -- movetime_ms is meant to
-        # bound the total time spent choosing this move, same contract as choose_move's `go movetime`.
         deadline = time.monotonic() + movetime_ms / 1000.0
         scored = []
         for mv, _ in ranked:
             child = searcher._expand(board, node, mv)
-            score = -searcher._negamax(board, child, searcher.max_depth - 1, -float("inf"), float("inf"), deadline)
+            score = -searcher._negamax(
+                board, child, searcher.max_depth - 1, -float("inf"), float("inf"), deadline
+            )
             board.pop()
             scored.append((mv, score))
         scored.sort(key=lambda x: -x[1])
-
         if ply < temperature_plies:
             move = random.choice(scored[: min(3, len(scored))])[0]
         else:
             move = scored[0][0]
-
-        side_to_move = board.turn  # capture before _expand pushes the move
+        side_to_move = board.turn
         from_to_id, promo_id = move_to_ids(move)
         soft_ids, soft_weights = scores_to_distribution(scored)
         window_from = hist_from[-MAX_PLIES:]
         window_to = hist_to[-MAX_PLIES:]
         window_promo = hist_promo[-MAX_PLIES:]
-        ply_records.append((window_from, window_to, window_promo, from_to_id, promo_id, side_to_move, soft_ids, soft_weights))
-
-        node = searcher._expand(board, node, move)  # advances board+state for real (no pop -- keeping this move)
+        ply_records.append(
+            (
+                window_from,
+                window_to,
+                window_promo,
+                from_to_id,
+                promo_id,
+                side_to_move,
+                soft_ids,
+                soft_weights,
+            )
+        )
+        node = searcher._expand(board, node, move)
         hist_from.append(move.from_square)
         hist_to.append(move.to_square)
         hist_promo.append(promo_id)
-
     result = board.result() if board.is_game_over() else "1/2-1/2"
     from data_pipeline import result_value
 
     examples = [
         (wf, wt, wp, tft, tp, result_value(result, side), sids, sw)
-        for (wf, wt, wp, tft, tp, side, sids, sw) in ply_records
+        for wf, wt, wp, tft, tp, side, sids, sw in ply_records
     ]
-    return examples, result
+    return (examples, result)
 
 
 def mirror_square(sq: int) -> int:
-    """Vertical rank-mirror (rank r -> 7-r, file unchanged) = XOR 56, since python-chess
-    squares are rank*8+file and XORing the rank's 3 bits with 0b111 gives 7-rank."""
     return sq ^ 56
 
 
@@ -119,13 +98,6 @@ def mirror_from_to(from_to_id: int) -> int:
 
 
 def mirror_example(ex):
-    """Chess has a real symmetry: mirror the board vertically and it's an equally legal
-    position from the other side's perspective. Doubles self-play data for free -- no
-    extra search needed, since it's a deterministic transform of an already-played game.
-    Piece types (promo ids), the value target (already side-to-move-relative, not
-    color-absolute), and soft-target weights are all unaffected by the mirror; only
-    squares need transforming. Padding slots (-1) in soft_ids are left untouched -- XORing
-    a sentinel would produce a bogus square id instead of staying a no-op padding marker."""
     wf, wt, wp, tft, tp, value, soft_ids, soft_w = ex
     m_wf = [mirror_square(s) for s in wf]
     m_wt = [mirror_square(s) for s in wt]
@@ -141,16 +113,15 @@ def main():
     ap.add_argument("--games", type=int, default=2000)
     ap.add_argument("--movetime", type=int, default=200)
     ap.add_argument("--temperature-plies", type=int, default=10)
-    ap.add_argument("--no-mirror-augment", action="store_true", help="disable color-flip data doubling")
+    ap.add_argument(
+        "--no-mirror-augment", action="store_true", help="disable color-flip data doubling"
+    )
     args = ap.parse_args()
-
     searcher = Searcher(checkpoint_path=args.checkpoint)
     if searcher.model is None:
         raise RuntimeError(f"no checkpoint found at {args.checkpoint} -- pretrain first")
-
     writer = ShardWriter(args.out_dir)
     results = {"1-0": 0, "0-1": 0, "1/2-1/2": 0}
-
     for g in range(args.games):
         examples, result = play_selfplay_game(searcher, args.movetime, args.temperature_plies)
         for ex in examples:
@@ -160,7 +131,6 @@ def main():
         results[result] = results.get(result, 0) + 1
         if (g + 1) % 20 == 0:
             print(f"[selfplay] {g + 1}/{args.games} games, results so far: {results}")
-
     writer.flush()
     print(f"[selfplay] DONE {args.games} games, results: {results}")