path: root/annotate_pgns/nnue_train.py

#!/usr/bin/env python3
"""Phase A training loop: local CPU, PyTorch, over the cached CSR
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
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 16
"""
import argparse
import json
import time
from pathlib import Path

import numpy as np
import torch

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")
        # 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


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 _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):
        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]
        material = data.material[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()),
            torch.from_numpy(material.astype(np.float32, copy=True)),
        )


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_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, material in iterate_batches(
                data, batch_size, shuffle=False, 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), 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_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():
    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("--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()

    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)
    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().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}")

    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, 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()
                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 += wdl_loss.item() * bs
                epoch_n += bs
                global_step += 1

                if global_step % args.eval_every_steps == 0:
                    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_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_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_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_wdl:.5f} "
                  f"holdout_material_mse={holdout_material:.4f} ({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}")

            # 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__":
    main()