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, 196 insertions, 25 deletions

diff --git a/annotate_pgns/nnue_train.py b/annotate_pgns/nnue_train.py
index 6112ab3..ea3d0a9 100644
--- a/annotate_pgns/nnue_train.py
+++ b/annotate_pgns/nnue_train.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """Phase A training loop: local CPU, PyTorch, over the cached CSR
-tensors build_training_tensors.py produces. Goal is pipeline/
+stensors 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
@@ -17,7 +17,7 @@ Usage:
         --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
+        --epochs 20 --batch-size 8192 --threads 16
 """
 import argparse
 import json
@@ -26,7 +26,6 @@ from pathlib import Path
 
 import numpy as np
 import torch
-import torch.nn as nn
 
 from nnue_model import NnueNet, DEFAULT_K
 
@@ -43,10 +42,17 @@ class CsrTensors:
         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")
+        # material.npy added 2026-10-17 (see nnue_model.py's
+        # material_probe) -- optional for backward compat with
+        # tensor dirs built before this field existed.
+        material_path = tensor_dir / "material.npy"
+        self.material = (np.load(material_path) if material_path.exists()
+                          else np.zeros_like(self.target))
         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
+        assert len(self.material) == self.n
 
     def __len__(self):
         return self.n
@@ -68,6 +74,16 @@ def _gather_batch(idx_arr, off_arr, row_ids):
     return flat, batch_off
 
 
+def _to_device(batch, device):
+    """Moves one iterate_batches() tuple onto device. A no-op (same
+    tensors returned) when device is "cpu", since CsrTensors already
+    hands back CPU tensors directly -- only a CUDA run pays this
+    transfer cost."""
+    if device == "cpu":
+        return batch
+    return tuple(t.to(device, non_blocking=True) for t in batch)
+
+
 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):
@@ -77,6 +93,7 @@ def iterate_batches(data: CsrTensors, batch_size, shuffle, rng):
         b_idx, b_off = _gather_batch(data.black_idx, data.black_off, rows_sorted)
         stm = data.stm[rows_sorted]
         target = data.target[rows_sorted]
+        material = data.material[rows_sorted]
         yield (
             torch.from_numpy(w_idx.astype(np.int64, copy=False)),
             torch.from_numpy(w_off),
@@ -84,23 +101,51 @@ def iterate_batches(data: CsrTensors, batch_size, shuffle, rng):
             torch.from_numpy(b_off),
             torch.from_numpy(stm.astype(np.int64, copy=False)),
             torch.from_numpy(target.copy()),
+            torch.from_numpy(material.astype(np.float32, copy=True)),
         )
 
 
-def evaluate(model, data, batch_size, k):
+def powered_loss(pred, target, exponent, reduction="mean"):
+    """|pred - target|^exponent, exponent=2 is plain MSE. Stockfish's
+    own training notes (nnue-pytorch wiki) report good results with
+    2.6 -- widens the loss ratio between small and large WDL-space
+    errors versus plain MSE, pushing the optimizer harder to fix
+    decisively-wrong (e.g. material-blind) predictions relative to
+    subtle positional ones. 2026-10-17: tried after plain MSE (exp=2)
+    showed near-identical score drops whether a pawn or a queen was
+    removed from a position (see sanity_check_net.py)."""
+    err = (pred - target).abs().pow(exponent)
+    return err.sum() if reduction == "sum" else err.mean()
+
+
+def evaluate(model, data, batch_size, k, loss_exponent, device="cpu"):
+    """Returns (wdl_loss, material_mse) -- kept separate rather than
+    pre-combined so the log/checkpoint can show whether the material
+    probe is actually converging, independently of whatever weight
+    the caller gives it in the training loss."""
     model.eval()
-    total_loss = 0.0
+    total_wdl_loss = 0.0
+    total_material_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(
+        for w_idx, w_off, b_idx, b_off, stm, target, material 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()
+            w_idx, w_off, b_idx, b_off, stm, target, material = _to_device(
+                (w_idx, w_off, b_idx, b_off, stm, target, material), device)
+            raw, material_pred = model.forward_with_material(w_idx, w_off, b_idx, b_off, stm)
+            pred = torch.sigmoid(raw / k)
+            wdl_loss = powered_loss(pred, target, loss_exponent, reduction="sum")
+            material_loss = nn_mse_sum(material_pred, material)
+            total_wdl_loss += wdl_loss.item()
+            total_material_loss += material_loss.item()
             total_n += len(target)
     model.train()
-    return total_loss / max(total_n, 1)
+    return total_wdl_loss / max(total_n, 1), total_material_loss / max(total_n, 1)
+
+
+def nn_mse_sum(pred, target):
+    return (pred - target).pow(2).sum()
 
 
 def main():
@@ -112,16 +157,65 @@ def main():
     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("--loss-exponent", type=float, default=2.0,
+                     help="exponent on |pred-target| in the loss (2.0 = "
+                          "plain MSE). Stockfish's own training notes "
+                          "report good results around 2.6 -- widens the "
+                          "loss ratio between small and large WDL-space "
+                          "errors, pushing harder to fix decisively-wrong "
+                          "predictions. See powered_loss()'s docstring.")
+    ap.add_argument("--material-loss-weight", type=float, default=0.0,
+                     help="weight on the auxiliary material_probe MSE "
+                          "term (plain pawns^2, NOT WDL-space -- so this "
+                          "is typically orders of magnitude larger in raw "
+                          "scale than the WDL loss; start small, e.g. "
+                          "0.001-0.01, and check sanity_check_net.py "
+                          "rather than assume a value works). 0.0 "
+                          "(default) disables it entirely -- the probe "
+                          "head still exists but gets no gradient, "
+                          "equivalent to not having added it.")
+    ap.add_argument("--threads", type=int, default=16,
+                     help="torch CPU thread count -- kept below this box's "
+                          "24 cores (not maxed out) since this is a shared "
+                          "multi-user box, matching the --workers default "
+                          "used elsewhere in this pipeline. Ignored when "
+                          "--device is not cpu.")
+    ap.add_argument("--device", default="cpu",
+                     help="torch device, e.g. 'cpu' (default, local "
+                          "iteration -- see this module's docstring), "
+                          "'cuda' (a rented GPU box, Phase B), or 'mps' "
+                          "(Apple Silicon). 'cuda'/'mps' with no matching "
+                          "torch/driver support available fails fast at "
+                          "startup rather than silently falling back to "
+                          "cpu.")
     ap.add_argument("--eval-every-steps", type=int, default=2000)
+    ap.add_argument("--patience", type=int, default=0,
+                     help="stop early if holdout_mse (the end-of-epoch "
+                          "value, not the more frequent mid-epoch one) "
+                          "fails to improve for this many consecutive "
+                          "epochs. 0 (default) disables early stopping -- "
+                          "always runs the full --epochs. Overfitting is "
+                          "still tracked either way (see best.json);  "
+                          "this only controls whether training keeps "
+                          "burning wall-clock time after it starts.")
     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)
+    if args.device.startswith("cuda") and not torch.cuda.is_available():
+        raise SystemExit(f"--device {args.device!r} requested but "
+                          f"torch.cuda.is_available() is False -- check the "
+                          f"CUDA driver/torch build on this box rather than "
+                          f"silently training on CPU")
+    if args.device == "mps" and not torch.backends.mps.is_available():
+        raise SystemExit(f"--device 'mps' requested but "
+                          f"torch.backends.mps.is_available() is False -- "
+                          f"needs Apple Silicon + a torch build with MPS "
+                          f"support, rather than silently training on CPU")
+
+    if args.device == "cpu":
+        torch.set_num_threads(args.threads)
     torch.manual_seed(args.seed)
 
     outdir = Path(args.outdir)
@@ -133,15 +227,43 @@ def main():
     holdout_data = CsrTensors(args.holdout_dir)
     print(f"  train: {len(train_data)} positions, holdout: {len(holdout_data)} positions")
 
-    model = NnueNet()
+    model = NnueNet().to(args.device)
     optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
     start_epoch = 0
     global_step = 0
+    best_holdout_mse = float("inf")
+    best_epoch = None
+    epochs_without_improvement = 0
+
+    best_path = outdir / "best.json"
+    if best_path.exists():
+        # Resuming (or re-running against an outdir that already has
+        # history): load prior best-epoch state rather than restarting
+        # the "best so far" search from scratch, which would otherwise
+        # let a later, worse epoch overwrite best.json just because
+        # this process forgot what the real best was.
+        prior_best = json.loads(best_path.read_text())
+        best_holdout_mse = prior_best["holdout_mse"]
+        best_epoch = prior_best["epoch"]
 
     if args.resume:
+        # map_location="cpu" always on load, then .to(device) below --
+        # a checkpoint saved from a CUDA run must still be resumable on
+        # a CPU-only box (and vice versa) without the load itself
+        # depending on CUDA being available.
         ckpt = torch.load(args.resume, map_location="cpu", weights_only=True)
         model.load_state_dict(ckpt["model"])
+        model.to(args.device)
         optimizer.load_state_dict(ckpt["optimizer"])
+        if args.device != "cpu":
+            # Adam's per-parameter state (exp_avg/exp_avg_sq) loads as
+            # whatever device it was saved from (map_location="cpu"
+            # above) -- move it to match the model's new device, or
+            # the first optimizer.step() throws a device-mismatch error.
+            for state in optimizer.state.values():
+                for key, val in state.items():
+                    if torch.is_tensor(val):
+                        state[key] = val.to(args.device)
         start_epoch = ckpt["epoch"] + 1
         global_step = ckpt.get("global_step", 0)
         print(f"Resumed from {args.resume} at epoch {start_epoch}")
@@ -155,40 +277,53 @@ def main():
             epoch_n = 0
             t_epoch = time.time()
 
-            for w_idx, w_off, b_idx, b_off, stm, target in iterate_batches(
+            for w_idx, w_off, b_idx, b_off, stm, target, material in iterate_batches(
                     train_data, args.batch_size, shuffle=True, rng=rng):
+                w_idx, w_off, b_idx, b_off, stm, target, material = _to_device(
+                    (w_idx, w_off, b_idx, b_off, stm, target, material), args.device)
                 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)
+                raw, material_pred = model.forward_with_material(w_idx, w_off, b_idx, b_off, stm)
+                pred = torch.sigmoid(raw / args.k)
+                wdl_loss = powered_loss(pred, target, args.loss_exponent)
+                material_loss = nn_mse_sum(material_pred, material) / len(target)
+                loss = wdl_loss + args.material_loss_weight * material_loss
                 loss.backward()
                 optimizer.step()
 
                 bs = len(target)
-                epoch_loss += loss.item() * bs
+                epoch_loss += wdl_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)
+                    holdout_wdl, holdout_material = evaluate(
+                        model, holdout_data, args.batch_size, args.k,
+                        args.loss_exponent, args.device)
                     elapsed = time.time() - t_start
                     rec = {
                         "step": global_step, "epoch": epoch,
                         "train_mse_running": epoch_loss / max(epoch_n, 1),
-                        "holdout_mse": holdout_mse,
+                        "holdout_mse": holdout_wdl,
+                        "holdout_material_mse": holdout_material,
                         "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"holdout_mse={holdout_wdl:.5f} "
+                          f"holdout_material_mse={holdout_material:.4f} "
                           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)
+            holdout_wdl, holdout_material = evaluate(
+                model, holdout_data, args.batch_size, args.k,
+                args.loss_exponent, args.device)
+            holdout_mse = holdout_wdl
             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)")
+                  f"holdout_mse={holdout_wdl:.5f} "
+                  f"holdout_material_mse={holdout_material:.4f} ({epoch_elapsed:.0f}s)")
 
             ckpt_path = outdir / f"ckpt_epoch{epoch:03d}.pt"
             torch.save({
@@ -201,7 +336,43 @@ def main():
             }, ckpt_path)
             print(f"  wrote {ckpt_path}")
 
-    print(f"Training done in {time.time() - t_start:.0f}s")
+            # Track the best-holdout epoch regardless of --patience, so
+            # a run that does nothing special still records which
+            # epoch to actually export -- see run_nnue_pipeline.py's
+            # best_checkpoint(), which reads this instead of just
+            # taking the highest epoch number (previously: always the
+            # LAST epoch trained, overfitting or not).
+            if holdout_mse < best_holdout_mse:
+                best_holdout_mse = holdout_mse
+                best_epoch = epoch
+                epochs_without_improvement = 0
+                best_path.write_text(json.dumps({
+                    "epoch": epoch,
+                    "holdout_mse": holdout_mse,
+                    "ckpt": ckpt_path.name,
+                }))
+                print(f"  new best (holdout_mse={holdout_mse:.5f})")
+            else:
+                epochs_without_improvement += 1
+                print(f"  not an improvement over epoch {best_epoch}'s "
+                      f"holdout_mse={best_holdout_mse:.5f} "
+                      f"({epochs_without_improvement} epoch(s) without "
+                      f"improvement)")
+                if args.patience and epochs_without_improvement >= args.patience:
+                    print(f"Stopping early: no holdout improvement for "
+                          f"{epochs_without_improvement} epochs "
+                          f"(--patience {args.patience})")
+                    break
+
+    print(f"Training done in {time.time() - t_start:.0f}s "
+          f"(best epoch: {best_epoch}, holdout_mse={best_holdout_mse:.5f})")
+
+    # Completion marker for run_nnue_pipeline.py's run_stage() --
+    # NOT the final epoch's checkpoint file, which --patience early
+    # stopping may legitimately mean never gets written.
+    (outdir / "train_complete.marker").write_text(json.dumps({
+        "best_epoch": best_epoch, "best_holdout_mse": best_holdout_mse,
+    }))
 
 
 if __name__ == "__main__":