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, 191 insertions, 66 deletions

diff --git a/annotate_pgns/fit_k_and_label.py b/annotate_pgns/fit_k_and_label.py
index 9fd09b0..11b2928 100644
--- a/annotate_pgns/fit_k_and_label.py
+++ b/annotate_pgns/fit_k_and_label.py
@@ -19,6 +19,16 @@ holdout data would leak information into what's supposed to be an
 untouched generalization check). The same fitted K is then applied
 (never refit) to label holdout positions too.
 
+Both the calibration-pair collection (--train) and the labeling pass
+(--train/--holdout) are sharded across --workers processes via
+multiprocessing.Pool, each worker vectorizing its own chunk of games
+with numpy -- this file used to run ~80 full Python-loop passes over
+every pair during the ternary search (2 mse() evals x 40 iterations),
+which was the slowest step in the post-annotation pipeline relative to
+how little actual compute it does. Per-game-chunk numpy vectorization
+plus Pool.imap (order-preserving, like build_training_tensors.py's
+existing pattern) turns that into array ops.
+
 Usage:
     python3 fit_k_and_label.py --train runs/sd10_twic_2M/games.train.jsonl \
         --holdout runs/sd10_twic_2M/games.holdout.jsonl \
@@ -27,7 +37,9 @@ Usage:
 """
 import argparse
 import json
-import math
+from multiprocessing import Pool
+
+import numpy as np
 
 
 def result_to_white_score(result):
@@ -35,77 +47,160 @@ def result_to_white_score(result):
     return {"1-0": 1.0, "0-1": 0.0, "1/2-1/2": 0.5}.get(result)
 
 
-def stm_relative_outcome(white_score, fen):
-    """Flip a White-perspective game result to the side-to-move's own
-    perspective at this FEN, matching the score field's own
-    side-to-move-relative convention."""
-    stm_is_white = fen.split()[1] == "w"
-    return white_score if stm_is_white else (1.0 - white_score)
-
-
-def collect_calibration_pairs(games_path):
-    """(score, r_outcome) for every non-exact position across all
-    games in this file."""
-    pairs = []
-    n_unknown_result = 0
-    with open(games_path) as f:
-        for line in f:
-            g = json.loads(line)
-            white_score = result_to_white_score(g["result"])
-            if white_score is None:
-                n_unknown_result += 1
-                continue
-            for pos in g["positions"]:
-                if pos.get("exact"):
-                    continue
-                r = stm_relative_outcome(white_score, pos["fen"])
-                pairs.append((pos["score"], r))
-    return pairs, n_unknown_result
+def sigmoid(scores, k):
+    return 1.0 / (1.0 + np.exp(-scores / k))
 
 
-def mse(pairs, k):
-    err = 0.0
-    for score, r in pairs:
-        p = 1.0 / (1.0 + math.exp(-score / k))
-        err += (r - p) ** 2
-    return err / max(len(pairs), 1)
+def mse(scores, outcomes, k):
+    p = sigmoid(scores, k)
+    return float(np.mean((outcomes - p) ** 2)) if len(scores) else 0.0
 
 
-def fit_k(pairs, lo=0.05, hi=20.0, iterations=40):
+def fit_k(scores, outcomes, lo=0.05, hi=20.0, iterations=40):
     for _ in range(iterations):
         m1 = lo + (hi - lo) / 3
         m2 = hi - (hi - lo) / 3
-        if mse(pairs, m1) < mse(pairs, m2):
+        if mse(scores, outcomes, m1) < mse(scores, outcomes, m2):
             hi = m2
         else:
             lo = m1
     return (lo + hi) / 2
 
 
-def label_file(in_path, out_path, k):
+def iter_chunks(path, chunk_size):
+    """Batches raw lines (one game per line) of an open jsonl file into
+    lists of up to chunk_size lines -- the unit of work handed to one
+    worker process at a time. Text, not parsed JSON, so the parse
+    itself happens in the worker, same division of labor as
+    build_training_tensors.py's iter_chunks."""
+    chunk = []
+    with open(path) as f:
+        for line in f:
+            chunk.append(line)
+            if len(chunk) >= chunk_size:
+                yield chunk
+                chunk = []
+    if chunk:
+        yield chunk
+
+
+def _non_exact_scores_and_outcomes(game):
+    """(scores, outcomes) numpy arrays for one game's non-exact
+    positions, or None if the game's result is unknown or it has no
+    non-exact positions."""
+    white_score = result_to_white_score(game["result"])
+    if white_score is None:
+        return None
+    non_exact = [p for p in game["positions"] if not p.get("exact")]
+    if not non_exact:
+        return None
+    scores = np.fromiter((p["score"] for p in non_exact), dtype=np.float64,
+                          count=len(non_exact))
+    is_white = np.fromiter(
+        (p["fen"].split()[1] == "w" for p in non_exact), dtype=bool,
+        count=len(non_exact))
+    outcomes = np.where(is_white, white_score, 1.0 - white_score)
+    return scores, outcomes
+
+
+def collect_chunk(lines):
+    """Runs in a worker process: parses a chunk of game lines and
+    returns (scores, outcomes, n_unknown_result) for every non-exact
+    position across those games."""
+    score_pieces = []
+    outcome_pieces = []
+    n_unknown_result = 0
+    for line in lines:
+        g = json.loads(line)
+        result = _non_exact_scores_and_outcomes(g)
+        if result is None:
+            if result_to_white_score(g["result"]) is None:
+                n_unknown_result += 1
+            continue
+        scores, outcomes = result
+        score_pieces.append(scores)
+        outcome_pieces.append(outcomes)
+    if not score_pieces:
+        return np.empty(0), np.empty(0), n_unknown_result
+    return (np.concatenate(score_pieces), np.concatenate(outcome_pieces),
+            n_unknown_result)
+
+
+def collect_calibration_pairs(games_path, workers, chunk_size):
+    """(scores, outcomes) numpy arrays for every non-exact position
+    across all games in this file, built by sharding games across
+    --workers processes."""
+    score_pieces = []
+    outcome_pieces = []
+    n_unknown_result = 0
+    chunks = iter_chunks(games_path, chunk_size)
+    with Pool(processes=workers) as pool:
+        for scores, outcomes, n_unk in pool.imap(collect_chunk, chunks):
+            if len(scores):
+                score_pieces.append(scores)
+                outcome_pieces.append(outcomes)
+            n_unknown_result += n_unk
+    if not score_pieces:
+        return np.empty(0), np.empty(0), n_unknown_result
+    return (np.concatenate(score_pieces), np.concatenate(outcome_pieces),
+            n_unknown_result)
+
+
+def label_chunk(args):
+    """Runs in a worker process: labels every non-exact position in a
+    chunk of game lines (vectorized per game with numpy) and returns
+    the re-serialized lines plus counts, in the same order they came
+    in."""
+    lines, k = args
+    out_lines = []
     n_games = 0
     n_positions = 0
     n_exact_skipped = 0
     n_unknown_result = 0
-    with open(in_path) as fin, open(out_path, "w") as fout:
-        for line in fin:
-            g = json.loads(line)
-            white_score = result_to_white_score(g["result"])
-            if white_score is None:
-                n_unknown_result += 1
-                fout.write(json.dumps(g) + "\n")
-                continue
-            for pos in g["positions"]:
-                n_positions += 1
-                if pos.get("exact"):
-                    n_exact_skipped += 1
-                    continue
-                r = stm_relative_outcome(white_score, pos["fen"])
-                p_search = 1.0 / (1.0 + math.exp(-pos["score"] / k))
-                pos["p_search"] = p_search
-                pos["r_outcome"] = r
-            fout.write(json.dumps(g) + "\n")
-            n_games += 1
+    for line in lines:
+        g = json.loads(line)
+        white_score = result_to_white_score(g["result"])
+        positions = g["positions"]
+        n_positions += len(positions)
+        if white_score is None:
+            n_unknown_result += 1
+            out_lines.append(json.dumps(g))
+            continue
+
+        idxs = [i for i, p in enumerate(positions) if not p.get("exact")]
+        n_exact_skipped += len(positions) - len(idxs)
+        if idxs:
+            scores = np.fromiter((positions[i]["score"] for i in idxs),
+                                  dtype=np.float64, count=len(idxs))
+            is_white = np.fromiter(
+                (positions[i]["fen"].split()[1] == "w" for i in idxs),
+                dtype=bool, count=len(idxs))
+            outcomes = np.where(is_white, white_score, 1.0 - white_score)
+            p_search = sigmoid(scores, k)
+            for j, i in enumerate(idxs):
+                positions[i]["p_search"] = float(p_search[j])
+                positions[i]["r_outcome"] = float(outcomes[j])
+
+        out_lines.append(json.dumps(g))
+        n_games += 1
+    return out_lines, n_games, n_positions, n_exact_skipped, n_unknown_result
+
+
+def label_file(in_path, out_path, k, workers, chunk_size):
+    n_games = 0
+    n_positions = 0
+    n_exact_skipped = 0
+    n_unknown_result = 0
+    chunks = ((chunk, k) for chunk in iter_chunks(in_path, chunk_size))
+    with Pool(processes=workers) as pool, open(out_path, "w") as fout:
+        for (out_lines, chunk_games, chunk_positions, chunk_exact,
+             chunk_unknown) in pool.imap(label_chunk, chunks):
+            for line in out_lines:
+                fout.write(line + "\n")
+            n_games += chunk_games
+            n_positions += chunk_positions
+            n_exact_skipped += chunk_exact
+            n_unknown_result += chunk_unknown
     return n_games, n_positions, n_exact_skipped, n_unknown_result
 
 
@@ -115,26 +210,56 @@ def main():
     ap.add_argument("--holdout", required=True)
     ap.add_argument("--train-out", required=True)
     ap.add_argument("--holdout-out", required=True)
+    ap.add_argument("--workers", type=int, default=16,
+                     help="worker processes for both the calibration-pair "
+                          "collection and the labeling pass (default 16, "
+                          "kept below this box's 24 cores to leave other "
+                          "jobs some breathing room)")
+    ap.add_argument("--chunk-size", type=int, default=2000,
+                     help="games per unit of work handed to a worker")
+    ap.add_argument("--fixed-k", type=float, default=None,
+                     help="use this K directly instead of re-fitting it "
+                          "from --train. 2026-10-16: needed when labeling "
+                          "a small supplementary slice (e.g. a self-play "
+                          "augmentation batch) to continue training an "
+                          "EXISTING checkpoint -- that checkpoint's "
+                          "weights were fit against labels built with its "
+                          "own K (nnue_model.py's DEFAULT_K unless noted "
+                          "otherwise), and re-fitting K from a small, "
+                          "distributionally skewed slice (e.g. one "
+                          "deliberately overweighted toward decisive-"
+                          "material positions) would silently relabel "
+                          "everything on a different, inconsistent scale "
+                          "-- the exact mismatch this module's own closing "
+                          "warning describes. Omit for a normal from-"
+                          "scratch fit on a large, representative pool.")
     args = ap.parse_args()
 
-    print(f"Collecting calibration pairs from {args.train} (TRAIN only) ...")
-    pairs, n_unknown = collect_calibration_pairs(args.train)
-    print(f"  {len(pairs)} (score, outcome) pairs "
-          f"({n_unknown} games skipped: unknown result)")
+    if args.fixed_k is not None:
+        k = args.fixed_k
+        print(f"Using fixed K = {k:.4f} (not re-fit from {args.train!r})")
+    else:
+        print(f"Collecting calibration pairs from {args.train} (TRAIN only) ...")
+        scores, outcomes, n_unknown = collect_calibration_pairs(
+            args.train, args.workers, args.chunk_size)
+        print(f"  {len(scores)} (score, outcome) pairs "
+              f"({n_unknown} games skipped: unknown result)")
 
-    print("Fitting K (ternary search, minimizing sigmoid-MSE against "
-          "actual game outcomes) ...")
-    k = fit_k(pairs)
-    final_mse = mse(pairs, k)
-    print(f"  K = {k:.2f}  (MSE at fit = {final_mse:.5f})")
+        print("Fitting K (ternary search, minimizing sigmoid-MSE against "
+              "actual game outcomes) ...")
+        k = fit_k(scores, outcomes)
+        final_mse = mse(scores, outcomes, k)
+        print(f"  K = {k:.2f}  (MSE at fit = {final_mse:.5f})")
 
     print(f"\nLabeling {args.train} -> {args.train_out} ...")
-    n_g, n_p, n_exact, n_unk = label_file(args.train, args.train_out, k)
+    n_g, n_p, n_exact, n_unk = label_file(
+        args.train, args.train_out, k, args.workers, args.chunk_size)
     print(f"  {n_g} games, {n_p} positions, {n_exact} exact (unlabeled), "
           f"{n_unk} unknown-result games (unlabeled)")
 
     print(f"Labeling {args.holdout} -> {args.holdout_out} (same K, not refit) ...")
-    n_g, n_p, n_exact, n_unk = label_file(args.holdout, args.holdout_out, k)
+    n_g, n_p, n_exact, n_unk = label_file(
+        args.holdout, args.holdout_out, k, args.workers, args.chunk_size)
     print(f"  {n_g} games, {n_p} positions, {n_exact} exact (unlabeled), "
           f"{n_unk} unknown-result games (unlabeled)")