Phase A NNUE: feature-indexing spec, annotation/training pipeline, dump nnuefeatures

Builds out the NNUE Phase A prototype per docs/NNUE.md:

- src/nnuefeatures.c: standalone HalfKAv2_hm feature-index reference
  implementation (new `dump nnuefeatures` command), sourced from
  Stockfish's half_ka_v2_hm.{h,cpp} (sf_16 tag), including the 0x88->64
  square conversion typhoon's own board representation needs.
- annotate_pgns/: sharded TWIC/self-play annotation harness
  (annotate.py), position extraction, quiet-position filtering (flat
  and game-shaped, plus an upstream --prefilter-quiet option),
  train/holdout game-level stratified split, WDL-sigmoid K fitting,
  the Python HalfKAv2_hm reference encoder (nnue_feature_spec.py) and
  its bit-for-bit cross-checker against the C side
  (cross_check_features.py), and a first PyTorch training pipeline
  (nnue_model.py, nnue_train.py, build_training_tensors.py) -- a
  256x2-32-32-1-ish net trained locally on CPU against ~1.85M quiet
  TWIC positions as a pipeline/framework validation pass.
- Finish an in-progress InitMagic/InitLMRTable/InitEval ->
  InitializeBitboardMagic/InitializeLMRTable/InitializeEval rename:
  chess.h's prototypes hadn't been updated to match the already-renamed
  definitions/call sites, leaving 4 implicit-function-declaration
  warnings.
- tests/arasan.epd: new curated EPD suite.
- docs/: add the project's design docs (previously untracked).

Co-Authored-By: Claude Sonnet 5 <[email protected]>

1 files changed, 208 insertions, 0 deletions

diff --git a/annotate_pgns/nnue_train.py b/annotate_pgns/nnue_train.py
new file mode 100644
index 0000000..6112ab3
--- /dev/null
+++ b/annotate_pgns/nnue_train.py
@@ -0,0 +1,208 @@
+#!/usr/bin/env python3
+"""Phase A training loop: local CPU, PyTorch, over the cached CSR
+tensors build_training_tensors.py produces. Goal is pipeline/
+framework validation (per NNUE.md) -- prove the data loader, model,
+loss, and checkpointing work end-to-end and get fast local iteration
+when something breaks, not to produce the best possible net. The
+real, larger Phase B run moves to a rented GPU against more data.
+
+CSR batching: each position is a variable-length "bag" of feature
+indices (3-32 of them). To build a batch for nn.EmbeddingBag without
+padding, gather the relevant CSR row-slices for a batch of position
+indices, concatenate them, and recompute offsets -- see
+_gather_batch() below.
+
+Usage:
+    python3 nnue_train.py \
+        --train-dir runs/sd10_twic_2M/tensors_train \
+        --holdout-dir runs/sd10_twic_2M/tensors_holdout \
+        --outdir runs/sd10_twic_2M/nnue_ckpt \
+        --epochs 20 --batch-size 8192 --threads 8
+"""
+import argparse
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from nnue_model import NnueNet, DEFAULT_K
+
+
+class CsrTensors:
+    """mmap-backed CSR feature tensors for one split (train or
+    holdout), as written by build_training_tensors.py."""
+
+    def __init__(self, tensor_dir):
+        tensor_dir = Path(tensor_dir)
+        self.white_idx = np.load(tensor_dir / "white_idx.npy", mmap_mode="r")
+        self.white_off = np.load(tensor_dir / "white_off.npy")
+        self.black_idx = np.load(tensor_dir / "black_idx.npy", mmap_mode="r")
+        self.black_off = np.load(tensor_dir / "black_off.npy")
+        self.stm = np.load(tensor_dir / "stm.npy")
+        self.target = np.load(tensor_dir / "target.npy")
+        self.n = len(self.target)
+        assert len(self.white_off) == self.n + 1
+        assert len(self.black_off) == self.n + 1
+        assert len(self.stm) == self.n
+
+    def __len__(self):
+        return self.n
+
+
+def _gather_batch(idx_arr, off_arr, row_ids):
+    """Build (flat_idx, batch_offsets) for an nn.EmbeddingBag call
+    covering exactly the bags named by row_ids (a 1D array of row
+    indices into off_arr). row_ids need not be sorted or contiguous."""
+    pieces = []
+    batch_off = np.empty(len(row_ids), dtype=np.int64)
+    cursor = 0
+    for i, r in enumerate(row_ids):
+        start, end = off_arr[r], off_arr[r + 1]
+        pieces.append(idx_arr[start:end])
+        batch_off[i] = cursor
+        cursor += end - start
+    flat = np.concatenate(pieces) if pieces else np.empty(0, dtype=np.int32)
+    return flat, batch_off
+
+
+def iterate_batches(data: CsrTensors, batch_size, shuffle, rng):
+    order = rng.permutation(data.n) if shuffle else np.arange(data.n)
+    for start in range(0, data.n, batch_size):
+        rows = order[start:start + batch_size]
+        rows_sorted = np.sort(rows)  # sequential mmap reads, not required but friendlier
+        w_idx, w_off = _gather_batch(data.white_idx, data.white_off, rows_sorted)
+        b_idx, b_off = _gather_batch(data.black_idx, data.black_off, rows_sorted)
+        stm = data.stm[rows_sorted]
+        target = data.target[rows_sorted]
+        yield (
+            torch.from_numpy(w_idx.astype(np.int64, copy=False)),
+            torch.from_numpy(w_off),
+            torch.from_numpy(b_idx.astype(np.int64, copy=False)),
+            torch.from_numpy(b_off),
+            torch.from_numpy(stm.astype(np.int64, copy=False)),
+            torch.from_numpy(target.copy()),
+        )
+
+
+def evaluate(model, data, batch_size, k):
+    model.eval()
+    total_loss = 0.0
+    total_n = 0
+    rng = np.random.default_rng(0)
+    with torch.no_grad():
+        for w_idx, w_off, b_idx, b_off, stm, target in iterate_batches(
+                data, batch_size, shuffle=False, rng=rng):
+            pred = model.win_probability(w_idx, w_off, b_idx, b_off, stm, k=k)
+            loss = nn.functional.mse_loss(pred, target, reduction="sum")
+            total_loss += loss.item()
+            total_n += len(target)
+    model.train()
+    return total_loss / max(total_n, 1)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--train-dir", required=True)
+    ap.add_argument("--holdout-dir", required=True)
+    ap.add_argument("--outdir", required=True)
+    ap.add_argument("--epochs", type=int, default=20)
+    ap.add_argument("--batch-size", type=int, default=8192)
+    ap.add_argument("--lr", type=float, default=1e-3)
+    ap.add_argument("--k", type=float, default=DEFAULT_K)
+    ap.add_argument("--threads", type=int, default=8,
+                     help="torch CPU thread count -- kept modest by default "
+                          "since this is a shared multi-user box")
+    ap.add_argument("--eval-every-steps", type=int, default=2000)
+    ap.add_argument("--resume", default=None,
+                     help="checkpoint .pt to resume from")
+    ap.add_argument("--seed", type=int, default=0)
+    args = ap.parse_args()
+
+    torch.set_num_threads(args.threads)
+    torch.manual_seed(args.seed)
+
+    outdir = Path(args.outdir)
+    outdir.mkdir(parents=True, exist_ok=True)
+    log_path = outdir / "train_log.jsonl"
+
+    print("Loading tensors ...")
+    train_data = CsrTensors(args.train_dir)
+    holdout_data = CsrTensors(args.holdout_dir)
+    print(f"  train: {len(train_data)} positions, holdout: {len(holdout_data)} positions")
+
+    model = NnueNet()
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+    start_epoch = 0
+    global_step = 0
+
+    if args.resume:
+        ckpt = torch.load(args.resume, map_location="cpu", weights_only=True)
+        model.load_state_dict(ckpt["model"])
+        optimizer.load_state_dict(ckpt["optimizer"])
+        start_epoch = ckpt["epoch"] + 1
+        global_step = ckpt.get("global_step", 0)
+        print(f"Resumed from {args.resume} at epoch {start_epoch}")
+
+    rng = np.random.default_rng(args.seed)
+    t_start = time.time()
+
+    with open(log_path, "a") as log_f:
+        for epoch in range(start_epoch, args.epochs):
+            epoch_loss = 0.0
+            epoch_n = 0
+            t_epoch = time.time()
+
+            for w_idx, w_off, b_idx, b_off, stm, target in iterate_batches(
+                    train_data, args.batch_size, shuffle=True, rng=rng):
+                optimizer.zero_grad()
+                pred = model.win_probability(w_idx, w_off, b_idx, b_off, stm, k=args.k)
+                loss = nn.functional.mse_loss(pred, target)
+                loss.backward()
+                optimizer.step()
+
+                bs = len(target)
+                epoch_loss += loss.item() * bs
+                epoch_n += bs
+                global_step += 1
+
+                if global_step % args.eval_every_steps == 0:
+                    holdout_mse = evaluate(model, holdout_data, args.batch_size, args.k)
+                    elapsed = time.time() - t_start
+                    rec = {
+                        "step": global_step, "epoch": epoch,
+                        "train_mse_running": epoch_loss / max(epoch_n, 1),
+                        "holdout_mse": holdout_mse,
+                        "elapsed_s": round(elapsed, 1),
+                        "positions_seen": epoch_n + epoch * len(train_data),
+                    }
+                    print(f"  step {global_step} epoch {epoch} "
+                          f"train_mse={rec['train_mse_running']:.5f} "
+                          f"holdout_mse={holdout_mse:.5f} "
+                          f"elapsed={elapsed:.0f}s")
+                    log_f.write(json.dumps(rec) + "\n")
+                    log_f.flush()
+
+            holdout_mse = evaluate(model, holdout_data, args.batch_size, args.k)
+            epoch_elapsed = time.time() - t_epoch
+            print(f"Epoch {epoch} done: train_mse={epoch_loss / max(epoch_n, 1):.5f} "
+                  f"holdout_mse={holdout_mse:.5f} ({epoch_elapsed:.0f}s)")
+
+            ckpt_path = outdir / f"ckpt_epoch{epoch:03d}.pt"
+            torch.save({
+                "model": model.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "epoch": epoch,
+                "global_step": global_step,
+                "holdout_mse": holdout_mse,
+                "args": vars(args),
+            }, ckpt_path)
+            print(f"  wrote {ckpt_path}")
+
+    print(f"Training done in {time.time() - t_start:.0f}s")
+
+
+if __name__ == "__main__":
+    main()