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, 326 insertions, 0 deletions

diff --git a/annotate_pgns/export_nnue_weights.py b/annotate_pgns/export_nnue_weights.py
new file mode 100644
index 0000000..4c8cb99
--- /dev/null
+++ b/annotate_pgns/export_nnue_weights.py
@@ -0,0 +1,326 @@
+#!/usr/bin/env python3
+"""Exports a trained nnue_model.py checkpoint into a flat binary file
+the C engine (src/nnueeval.c's NnueLoadWeights) can read directly --
+no parsing, just a fixed-order sequence of arrays behind a small
+header.
+
+File format (all little-endian), version "NNU3":
+    magic:     4 bytes, ASCII "NNU3"
+    dims:      4x uint32 -- NNUE_DIMENSIONS, L1, L2, L3 (sanity-checked
+               against the C side's own compile-time constants on
+               load, so a mismatched build can't silently misread the
+               file)
+    ft_scale:  1x float32 -- fixed-point scale for ft.weight/ft_bias
+               (int16). NnueForward divides the int32 accumulator sum
+               by this to get the float "fc1 input" value.
+    act_scale: 3x float32 (AS1, AS2, AS3) -- fixed-point scales used
+               to quantize each stage's post-ReLU activation to uint8
+               before feeding the next layer's int8 dot product:
+                 AS1 quantizes fc1's input (the ReLU'd, concatenated
+                     accumulator halves)
+                 AS2 quantizes fc1's ReLU'd output (fc2's input)
+                 AS3 quantizes fc2's ReLU'd output (fc3's input)
+               Derived from real sampled activations (see
+               CalibrateActivationScales) with headroom, NOT from a
+               theoretical max -- see docs/NNUE.md's "Quantization
+               plan" / the int8 float-vs-quantized correctness check.
+    w_scale:   3x float32 (WS1, WS2, WS3) -- fixed-point scales for
+               fc1.weight/fc2.weight/fc3.weight (int8), one per layer,
+               each derived from that layer's own max|.| (exact, not
+               sampled -- weights are fixed, known precisely from the
+               checkpoint).
+    then, in this exact order:
+        ft.weight   [NNUE_DIMENSIONS x L1]   int16  (quantized @ ft_scale)
+        ft_bias     [L1]                     int16  (quantized @ ft_scale)
+        fc1.weight  [L2 x (2*L1)]            int8   (quantized @ WS1)
+        fc1.bias    [L2]                     float32 (unquantized --
+                                              added post-dequant, tiny
+                                              array, not worth it)
+        fc2.weight  [L3 x L2]                int8   (quantized @ WS2)
+        fc2.bias    [L3]                     float32
+        fc3.weight  [1 x L3]                 int8   (quantized @ WS3)
+        fc3.bias    [1]                      float32
+
+Both the embedding table (int16) and the FC head (int8) are quantized
+now -- see docs/NNUE.md's "Quantization plan". Measured accuracy cost
+of the FC-head int8 step, 5000 real holdout positions vs the exact
+float reference: mean |delta| 1.55cp, median 0.99cp, max 20.91cp, zero
+clipping -- an order of magnitude below the sd10-vs-sd16 search-depth
+noise (mean 19cp) the project already tolerates in its training
+labels. See the conversation/commit history for the simulation that
+produced those numbers; CalibrateActivationScales below reproduces the
+scale derivation deterministically from a real checkpoint + its
+training tensors.
+
+Usage:
+    python3 export_nnue_weights.py \
+        runs/sd10_twic_2M/nnue_ckpt/ckpt_epoch019.pt \
+        runs/sd10_twic_2M/nnue_ckpt/weights.nnue \
+        --tensor-dir runs/sd10_twic_2M/tensors_train
+"""
+import argparse
+import struct
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from nnue_feature_spec import DIMENSIONS, PS_NB
+from nnue_model import L1, L2, L3, NnueNet
+
+# Fixed-point scale for ft.weight/ft_bias. Checked against
+# ckpt_epoch019.pt (ft.weight max|.|=5.01, ft_bias max|.|=0.52): at
+# this scale, 5.01*4096 ~= 20502, comfortably under int16's 32767
+# ceiling (~1.6x headroom for a larger Phase B net before this needs
+# re-deriving). See docs/NNUE.md's "Quantization plan".
+FT_SCALE = 4096.0
+
+# Headroom multiplier applied to the largest value seen in a real
+# sample before deriving each activation scale -- protects against
+# rare tail values in the full corpus that a finite sample might not
+# have hit. 1.3x chosen as a round, conservative margin; revisit if a
+# larger/differently-distributed Phase B net ever reports clipping.
+ACT_HEADROOM = 1.3
+
+
+def CalibrateActivationScales(model, tensor_dir, n_sample=50000, seed=42):
+    """Runs n_sample real positions from tensor_dir (one of
+    build_training_tensors.py's output directories) through the float
+    model, capturing each stage's real post-ReLU activation
+    distribution via forward hooks, and returns (AS1, AS2, AS3) --
+    uint8 fixed-point scales with ACT_HEADROOM margin over the
+    observed max. Must run against the SAME checkpoint being
+    exported -- these scales are not portable across checkpoints."""
+    d = Path(tensor_dir)
+    white_idx = np.load(d / "white_idx.npy"); white_off = np.load(d / "white_off.npy")
+    black_idx = np.load(d / "black_idx.npy"); black_off = np.load(d / "black_off.npy")
+    stm = np.load(d / "stm.npy")
+    n = len(stm)
+    rng = np.random.default_rng(seed)
+    rows = rng.choice(n, size=min(n_sample, n), replace=False)
+    rows.sort()
+
+    captured = {}
+    h1 = model.fc1.register_forward_hook(
+        lambda m, i, o: captured.setdefault("fc1_in", []).append(i[0].detach()))
+    h2 = model.fc2.register_forward_hook(
+        lambda m, i, o: captured.setdefault("fc2_in", []).append(i[0].detach()))
+    h3 = model.fc3.register_forward_hook(
+        lambda m, i, o: captured.setdefault("fc3_in", []).append(i[0].detach()))
+
+    with torch.no_grad():
+        for i in range(0, len(rows), 5000):
+            chunk = rows[i:i + 5000]
+            w_idx_list, w_off_list, b_idx_list, b_off_list = [], [], [], []
+            off_w, off_b = 0, 0
+            for r in chunk:
+                wi = white_idx[white_off[r]:white_off[r + 1]]
+                bi = black_idx[black_off[r]:black_off[r + 1]]
+                w_idx_list.append(wi); w_off_list.append(off_w); off_w += len(wi)
+                b_idx_list.append(bi); b_off_list.append(off_b); off_b += len(bi)
+            w_idx = torch.from_numpy(np.concatenate(w_idx_list).astype(np.int64))
+            w_off = torch.tensor(w_off_list, dtype=torch.int64)
+            b_idx = torch.from_numpy(np.concatenate(b_idx_list).astype(np.int64))
+            b_off = torch.tensor(b_off_list, dtype=torch.int64)
+            stm_t = torch.from_numpy(stm[chunk].astype(np.int64))
+            model(w_idx, w_off, b_idx, b_off, stm_t)
+
+    h1.remove(); h2.remove(); h3.remove()
+
+    scales = []
+    for key in ("fc1_in", "fc2_in", "fc3_in"):
+        t = torch.cat(captured[key]).float()
+        m = t.max().item()
+        scales.append(255.0 / (m * ACT_HEADROOM))
+    return tuple(scales)
+
+
+def QuantizeFt(arr):
+    """Rounds + clamps a float32 array to int16 at FT_SCALE. Returns
+    (int16 ndarray, number of values that clipped)."""
+    scaled = np.round(arr.astype(np.float64) * FT_SCALE)
+    clipped = int(np.count_nonzero((scaled < -32768) | (scaled > 32767)))
+    scaled = np.clip(scaled, -32768, 32767)
+    return scaled.astype(np.int16), clipped
+
+
+def QuantizeWeightS8(arr, scale):
+    """Rounds + clamps a float32 array to int8 at the given scale.
+    Returns (int8 ndarray, number of values that clipped)."""
+    scaled = np.round(arr.astype(np.float64) * scale)
+    clipped = int(np.count_nonzero((scaled < -128) | (scaled > 127)))
+    scaled = np.clip(scaled, -128, 127)
+    return scaled.astype(np.int8), clipped
+
+
+def CoalesceFactorization(sd, tensor_dir, n_check=4096, seed=123):
+    """If the checkpoint has a trained ft_p factorization table (see
+    nnue_model.py's 2026-10-18 note), folds it into ft.weight in
+    place (idx % PS_NB is the P-factor's row for real feature idx)
+    and returns a modified copy of sd with ft_p dropped -- the
+    coalesced ft alone must reproduce the factored forward pass, so
+    ft_p carries zero inference-time cost. No-op (returns sd
+    unchanged) if ft_p isn't present (an older, unfactored
+    checkpoint).
+
+    Required correctness check, per NNUE.md: a bad coalesce silently
+    produces a WORSE net (real rows learn less because the factor
+    absorbed shared signal) with no sign of it in the training loss
+    curve. Verified here by comparing the pre-coalesce (factored,
+    two-table) forward pass against the post-coalesce (plain,
+    ft-only) forward pass on n_check real positions sampled from
+    tensor_dir -- must agree to ~1e-5 or this refuses to export."""
+    if "ft_p.weight" not in sd:
+        return sd
+
+    ft_p_weight = sd["ft_p.weight"]
+    if tuple(ft_p_weight.shape) != (PS_NB, L1):
+        sys.exit(f"ft_p.weight: expected shape {(PS_NB, L1)}, checkpoint "
+                  f"has {tuple(ft_p_weight.shape)} -- wrong model version?")
+
+    idx_mod = torch.arange(DIMENSIONS) % PS_NB
+    coalesced_ft_weight = sd["ft.weight"] + ft_p_weight[idx_mod]
+
+    sd_factored = dict(sd)
+    model_factored = NnueNet()
+    model_factored.load_state_dict(sd_factored)
+    model_factored.eval()
+
+    sd_plain = dict(sd)
+    sd_plain["ft.weight"] = coalesced_ft_weight
+    sd_plain["ft_p.weight"] = torch.zeros_like(ft_p_weight)
+    model_plain = NnueNet()
+    model_plain.load_state_dict(sd_plain)
+    model_plain.eval()
+
+    d = Path(tensor_dir)
+    white_idx = np.load(d / "white_idx.npy"); white_off = np.load(d / "white_off.npy")
+    black_idx = np.load(d / "black_idx.npy"); black_off = np.load(d / "black_off.npy")
+    stm = np.load(d / "stm.npy")
+    rng = np.random.default_rng(seed)
+    rows = rng.choice(len(stm), size=min(n_check, len(stm)), replace=False)
+    rows.sort()
+
+    w_idx_list, w_off_list, b_idx_list, b_off_list = [], [], [], []
+    off_w, off_b = 0, 0
+    for r in rows:
+        wi = white_idx[white_off[r]:white_off[r + 1]]
+        bi = black_idx[black_off[r]:black_off[r + 1]]
+        w_idx_list.append(wi); w_off_list.append(off_w); off_w += len(wi)
+        b_idx_list.append(bi); b_off_list.append(off_b); off_b += len(bi)
+    w_idx = torch.from_numpy(np.concatenate(w_idx_list).astype(np.int64))
+    w_off = torch.tensor(w_off_list, dtype=torch.int64)
+    b_idx = torch.from_numpy(np.concatenate(b_idx_list).astype(np.int64))
+    b_off = torch.tensor(b_off_list, dtype=torch.int64)
+    stm_t = torch.from_numpy(stm[rows].astype(np.int64))
+
+    with torch.no_grad():
+        out_factored = model_factored(w_idx, w_off, b_idx, b_off, stm_t)
+        out_plain = model_plain(w_idx, w_off, b_idx, b_off, stm_t)
+    max_delta = (out_factored - out_plain).abs().max().item()
+    if max_delta > 1e-4:
+        sys.exit(f"Coalesce check FAILED: max |factored - coalesced| = "
+                  f"{max_delta:.6g} over {n_check} real positions (expected "
+                  f"~1e-5) -- refusing to export a bad coalesce. Check "
+                  f"ft_p's shape/init and the idx % PS_NB formula.")
+    print(f"Coalesce check OK: max |factored - coalesced| = {max_delta:.3g} "
+          f"over {n_check} real positions")
+
+    sd_out = dict(sd)
+    sd_out["ft.weight"] = coalesced_ft_weight
+    del sd_out["ft_p.weight"]
+    return sd_out
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("checkpoint")
+    ap.add_argument("outfile")
+    ap.add_argument("--tensor-dir", required=True,
+                     help="a build_training_tensors.py output dir, used "
+                          "to calibrate the FC head's activation scales "
+                          "against real positions (e.g. tensors_train)")
+    args = ap.parse_args()
+
+    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
+    sd = ckpt["model"]
+    sd = CoalesceFactorization(sd, args.tensor_dir)
+
+    expected_shapes = {
+        "ft.weight": (DIMENSIONS, L1),
+        "ft_bias": (L1,),
+        "fc1.weight": (L2, 2 * L1),
+        "fc1.bias": (L2,),
+        "fc2.weight": (L3, L2),
+        "fc2.bias": (L3,),
+        "fc3.weight": (1, L3),
+        "fc3.bias": (1,),
+    }
+    for name, shape in expected_shapes.items():
+        if tuple(sd[name].shape) != shape:
+            sys.exit(f"{name}: expected shape {shape}, checkpoint has "
+                      f"{tuple(sd[name].shape)} -- wrong model version?")
+
+    model = NnueNet()
+    # strict=False: CoalesceFactorization deliberately drops ft_p.weight
+    # (folded into ft.weight above) when the checkpoint had it, but
+    # NnueNet always declares the ft_p parameter regardless -- missing_keys
+    # is expected to be exactly ["ft_p.weight"] in that case, or empty for
+    # an older unfactored checkpoint CoalesceFactorization no-op'd on.
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    bad_missing = [k for k in missing if k != "ft_p.weight"]
+    if bad_missing or unexpected:
+        sys.exit(f"load_state_dict mismatch beyond the expected ft_p.weight "
+                  f"drop -- missing={bad_missing}, unexpected={unexpected}")
+    model.eval()
+
+    AS1, AS2, AS3 = CalibrateActivationScales(model, args.tensor_dir)
+
+    ft_weight_f32 = sd["ft.weight"].detach().contiguous().to(torch.float32).numpy()
+    ft_bias_f32 = sd["ft_bias"].detach().contiguous().to(torch.float32).numpy()
+    fc1_w_f32 = sd["fc1.weight"].detach().contiguous().to(torch.float32).numpy()
+    fc2_w_f32 = sd["fc2.weight"].detach().contiguous().to(torch.float32).numpy()
+    fc3_w_f32 = sd["fc3.weight"].detach().contiguous().to(torch.float32).numpy()
+
+    ft_weight_q, clipped_ftw = QuantizeFt(ft_weight_f32)
+    ft_bias_q, clipped_ftb = QuantizeFt(ft_bias_f32)
+
+    WS1 = 127.0 / np.abs(fc1_w_f32).max()
+    WS2 = 127.0 / np.abs(fc2_w_f32).max()
+    WS3 = 127.0 / np.abs(fc3_w_f32).max()
+
+    fc1_w_q, clipped_w1 = QuantizeWeightS8(fc1_w_f32, WS1)
+    fc2_w_q, clipped_w2 = QuantizeWeightS8(fc2_w_f32, WS2)
+    fc3_w_q, clipped_w3 = QuantizeWeightS8(fc3_w_f32, WS3)
+
+    total_clipped = (clipped_ftw + clipped_ftb + clipped_w1 + clipped_w2 +
+                      clipped_w3)
+    if total_clipped:
+        sys.exit(f"Quantize: {total_clipped} value(s) clipped -- "
+                  f"re-check headroom/scales before shipping this export "
+                  f"(see docs/NNUE.md)")
+
+    with open(args.outfile, "wb") as f:
+        f.write(b"NNU3")
+        f.write(struct.pack("<4I", DIMENSIONS, L1, L2, L3))
+        f.write(struct.pack("<f", FT_SCALE))
+        f.write(struct.pack("<3f", AS1, AS2, AS3))
+        f.write(struct.pack("<3f", WS1, WS2, WS3))
+        f.write(ft_weight_q.tobytes())
+        f.write(ft_bias_q.tobytes())
+        f.write(fc1_w_q.tobytes())
+        f.write(sd["fc1.bias"].detach().contiguous().to(torch.float32).numpy().tobytes())
+        f.write(fc2_w_q.tobytes())
+        f.write(sd["fc2.bias"].detach().contiguous().to(torch.float32).numpy().tobytes())
+        f.write(fc3_w_q.tobytes())
+        f.write(sd["fc3.bias"].detach().contiguous().to(torch.float32).numpy().tobytes())
+
+    print(f"Wrote {args.outfile} from {args.checkpoint}")
+    print(f"  ft_scale={FT_SCALE}")
+    print(f"  act_scales (AS1,AS2,AS3)=({AS1:.4f},{AS2:.4f},{AS3:.4f})")
+    print(f"  weight_scales (WS1,WS2,WS3)=({WS1:.4f},{WS2:.4f},{WS3:.4f})")
+
+
+if __name__ == "__main__":
+    main()