#!/usr/bin/env python3 """Phase A training loop: local CPU, PyTorch, over the cached CSR tensors 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 8 """ import argparse import json import time from pathlib import Path import numpy as np import torch import torch.nn as nn 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") 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 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 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] 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()), ) def evaluate(model, data, batch_size, k): model.eval() total_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( 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() total_n += len(target) model.train() return total_loss / max(total_n, 1) 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("--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("--eval-every-steps", type=int, default=2000) 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) 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() optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) start_epoch = 0 global_step = 0 if args.resume: ckpt = torch.load(args.resume, map_location="cpu", weights_only=True) model.load_state_dict(ckpt["model"]) optimizer.load_state_dict(ckpt["optimizer"]) 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 in iterate_batches( train_data, args.batch_size, shuffle=True, rng=rng): 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) loss.backward() optimizer.step() bs = len(target) epoch_loss += 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) elapsed = time.time() - t_start rec = { "step": global_step, "epoch": epoch, "train_mse_running": epoch_loss / max(epoch_n, 1), "holdout_mse": holdout_mse, "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"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) 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)") 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}") print(f"Training done in {time.time() - t_start:.0f}s") if __name__ == "__main__": main()