Wire NNUE eval into search, build the training pipeline, and add a fast

move/eval-agreement benchmark against Stockfish (tests/lop)

Implements the architecture decisions in src/docs/NNUE.md: HalfKAv2_hm
feature set, int16/int8 quantization, per-ply accumulator backup for
king moves, and a lazy-eval-ready integration into search.c/eval.c.
Adds the full annotate_pgns/ training pipeline (extraction, K-fit/
labeling, quiet filtering, tensor building, training, export) plus an
end-to-end orchestrator (run_nnue_pipeline.py) and validation tooling
(sanity_check_net.py, compare_vs_stockfish_net.py). Adds tests/lop:
a 10k-position list-of-positions benchmark comparing engine eval rank
and move choice against pre-computed Stockfish/lc0 references, much
cheaper than a full self-play match for day-to-day progress checks.

First sd10-matched self-play result against classic HCE: LLR +2.28
(H1=+2.94), -72=119+169 over 360 games, confirming the net is
"smarter per node" at matched depth. A real-time (--st 4) match is
currently weaker (candidate loses ~3 plies of depth to a ~3.3-3.7x
nps gap against an unrecalibrated search), consistent with
NNUE.md's adoption-criterion section -- pruning/reduction constants
still need retuning for the new eval's score scale before a fair
real-time comparison is possible.

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

1 files changed, 183 insertions, 0 deletions

diff --git a/annotate_pgns/compare_vs_stockfish_net.py b/annotate_pgns/compare_vs_stockfish_net.py
new file mode 100644
index 0000000..0fe7514
--- /dev/null
+++ b/annotate_pgns/compare_vs_stockfish_net.py
@@ -0,0 +1,183 @@
+#!/usr/bin/env python3
+"""Compares our trained checkpoints against a real, matching-commit
+Stockfish build's NNUE evaluation on sampled real positions.
+
+NOT a weight-reuse step (see conversation: the goal is a from-scratch
+net) -- this is purely a diagnostic to see how far our net's score
+scale/ranking diverges from a known-strong reference, and whether
+feature factorization measurably closes that gap vs. the unfactored
+run.
+
+Stockfish's "NNUE evaluation" line (not "Final evaluation", which
+blends in a classical/hybrid term we have no equivalent of) is the
+directly comparable quantity -- both it and our forward() are pure
+NNUE output, no classical blending, both white-side-relative after
+sign conversion.
+
+Usage:
+    python3 compare_vs_stockfish_net.py \
+        --sample-file runs/stockfish_selfplay/games.holdout.labeled.quiet.jsonl \
+        --n-positions 1000 \
+        --sf-binary /tmp/sf_38e616/src/stockfish \
+        --sf-net nn-fdc1d0fe6455.nnue \
+        --ckpt-unfactored runs/stockfish_selfplay/nnue_ckpt_lambda075/ckpt_epoch033.pt \
+        --ckpt-factored runs/stockfish_selfplay/nnue_ckpt_lambda075_factored/ckpt_epoch033.pt
+"""
+import argparse
+import json
+import random
+import re
+import subprocess
+import sys
+
+import chess
+import numpy as np
+import torch
+
+from nnue_feature_spec import active_feature_indices
+from nnue_model import NnueNet
+
+NNUE_EVAL_RE = re.compile(r"NNUE evaluation\s+([+-]?\d+\.\d+) \(white side\)")
+
+
+def sample_fens(sample_file, n, seed):
+    rng = random.Random(seed)
+    reservoir = []
+    with open(sample_file) as f:
+        seen = 0
+        for line in f:
+            game = json.loads(line)
+            for pos in game["positions"]:
+                if pos.get("exact"):
+                    continue
+                seen += 1
+                if len(reservoir) < n:
+                    reservoir.append(pos["fen"])
+                else:
+                    j = rng.randint(0, seen - 1)
+                    if j < n:
+                        reservoir[j] = pos["fen"]
+    return reservoir
+
+
+def load_model(ckpt_path):
+    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+    model = NnueNet()
+    missing, unexpected = model.load_state_dict(ckpt["model"], strict=False)
+    bad_missing = [k for k in missing if k != "ft_p.weight"]
+    if bad_missing or unexpected:
+        sys.exit(f"{ckpt_path}: load_state_dict mismatch -- "
+                  f"missing={bad_missing}, unexpected={unexpected}")
+    model.eval()
+    return model
+
+
+def our_score_white_relative(model, fen):
+    board = chess.Board(fen)
+    w = active_feature_indices(board, chess.WHITE)
+    b = active_feature_indices(board, chess.BLACK)
+    w_idx = torch.tensor(w, dtype=torch.long)
+    w_off = torch.tensor([0], dtype=torch.long)
+    b_idx = torch.tensor(b, dtype=torch.long)
+    b_off = torch.tensor([0], dtype=torch.long)
+    stm = torch.tensor([1 if board.turn == chess.WHITE else 0], dtype=torch.long)
+    with torch.no_grad():
+        raw = model(w_idx, w_off, b_idx, b_off, stm).item()
+    return raw if board.turn == chess.WHITE else -raw
+
+
+def stockfish_scores_white_relative(sf_binary, sf_net, fens):
+    """One persistent Stockfish process, one 'position fen ...; eval'
+    round-trip per FEN, parsing the 'NNUE evaluation' line. Much
+    cheaper than spawning a fresh process per position."""
+    proc = subprocess.Popen(
+        [sf_binary], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+        stderr=subprocess.DEVNULL, text=True, bufsize=1)
+
+    def send(cmd):
+        proc.stdin.write(cmd + "\n")
+        proc.stdin.flush()
+
+    send(f"setoption name EvalFile value {sf_net}")
+    send("isready")
+    while True:
+        line = proc.stdout.readline()
+        if not line or "readyok" in line:
+            break
+
+    scores = []
+    for fen in fens:
+        send(f"position fen {fen}")
+        send("eval")
+        score = None
+        while True:
+            line = proc.stdout.readline()
+            if not line:
+                sys.exit("Stockfish process died mid-eval -- net likely "
+                         "failed to load (check EvalFile path/error output)")
+            m = NNUE_EVAL_RE.search(line)
+            if m:
+                score = float(m.group(1))
+            if "Final evaluation" in line:
+                break
+        scores.append(score)
+    send("quit")
+    proc.wait(timeout=5)
+    return scores
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__,
+                                  formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--sample-file", required=True)
+    ap.add_argument("--n-positions", type=int, default=1000)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--sf-binary", required=True)
+    ap.add_argument("--sf-net", required=True)
+    ap.add_argument("--ckpt-unfactored", required=True)
+    ap.add_argument("--ckpt-factored", required=True)
+    args = ap.parse_args()
+
+    print(f"Sampling {args.n_positions} non-exact positions from {args.sample_file} ...")
+    fens = sample_fens(args.sample_file, args.n_positions, args.seed)
+    print(f"  got {len(fens)} FENs")
+
+    print("Loading checkpoints ...")
+    model_unfactored = load_model(args.ckpt_unfactored)
+    model_factored = load_model(args.ckpt_factored)
+
+    print(f"Running {args.sf_binary} (net={args.sf_net}) ...")
+    sf_scores = stockfish_scores_white_relative(args.sf_binary, args.sf_net, fens)
+
+    print("Running our checkpoints ...")
+    ours_unfactored = [our_score_white_relative(model_unfactored, fen) for fen in fens]
+    ours_factored = [our_score_white_relative(model_factored, fen) for fen in fens]
+
+    sf = np.array(sf_scores, dtype=np.float64)
+    uf = np.array(ours_unfactored, dtype=np.float64)
+    fa = np.array(ours_factored, dtype=np.float64)
+
+    def report(name, ours):
+        diff = ours - sf
+        corr = np.corrcoef(ours, sf)[0, 1]
+        # best-fit scale (least squares, no intercept) -- how far off is
+        # our net's overall SCALE from Stockfish's, independent of the
+        # per-position noise captured by correlation/residual.
+        scale = float(np.dot(ours, sf) / np.dot(ours, ours))
+        resid_after_scale = ours * scale - sf
+        print(f"\n--- {name} vs Stockfish (n={len(sf)}) ---")
+        print(f"  correlation:            {corr:.4f}")
+        print(f"  raw mean|diff|:         {np.abs(diff).mean():.3f} pawns")
+        print(f"  raw median diff:        {np.median(diff):.3f} pawns")
+        print(f"  best-fit scale (ours->sf): {scale:.3f}")
+        print(f"  mean|resid| after scale-fit: {np.abs(resid_after_scale).mean():.3f} pawns")
+
+    report("Unfactored", uf)
+    report("Factored", fa)
+
+    np.savez("/tmp/compare_vs_sf.npz", sf=sf, unfactored=uf, factored=fa, fens=np.array(fens))
+    print("\nSaved raw arrays to /tmp/compare_vs_sf.npz")
+
+
+if __name__ == "__main__":
+    sys.exit(main())