#!/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()