1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
|
#!/usr/bin/env python3
"""Fits the lazy-eval probe (docs/NNUE.md's "Lazy-evaluation tiering
plan"): a ridge regression on the trained net's own accumulator
(fc1's input -- the ReLU'd, concatenated 512-wide per-perspective
halves), targeting that SAME checkpoint's raw forward() output
(pre-calibration -- see export_nnue_weights.py's CalibrateScoreMapping;
NnueProbeScore applies the identical (CALIB_A, CALIB_B) correction
itself at runtime, so the probe should target raw, not the calibrated
cp value, to stay consistent with that).
Must be re-run for every new checkpoint -- these weights are fit
against one specific net's accumulator and are not portable across
generations (same rule as material/score calibration).
Writes the NPRB format NnueProbeLoadWeights (src/nnueeval.c) reads:
magic: "NPRB"
dim: uint32 (2*L1)
mu: dim x float32
sigma: dim x float32
weight: dim x float32
bias: 1x float32
Usage:
python3 fit_lazy_probe.py \
runs/sf_20261020_recal/nnue_ckpt/ckpt_epoch015.pt \
runs/sf_20261020_recal/nnue_probe.bin \
--train-tensor-dir runs/sf_20261020_recal/tensors_train \
--holdout-tensor-dir runs/sf_20261020_recal/tensors_holdout
"""
import argparse
import struct
from pathlib import Path
import numpy as np
import torch
from nnue_model import NnueNet
RIDGE_LAMBDA = 1.0
def _batched_accumulator_and_raw(model, tensor_dir, rows, batch_size=5000):
"""Returns (accum [N, 2*L1] float64, raw [N] float64) for the given
row indices into tensor_dir's CSR arrays."""
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")
accums, raws = [], []
with torch.no_grad():
for i in range(0, len(rows), batch_size):
chunk = rows[i:i + batch_size]
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))
acc = model.accumulator(w_idx, w_off, b_idx, b_off, stm_t)
raw = model.forward(w_idx, w_off, b_idx, b_off, stm_t)
accums.append(acc.numpy())
raws.append(raw.numpy())
return (np.concatenate(accums).astype(np.float64),
np.concatenate(raws).astype(np.float64))
def main():
ap = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("checkpoint")
ap.add_argument("outfile")
ap.add_argument("--train-tensor-dir", required=True,
help="fit the ridge regression against positions from here")
ap.add_argument("--holdout-tensor-dir", required=True,
help="report residual stats (not used for fitting) against here")
ap.add_argument("--n-fit", type=int, default=80000)
ap.add_argument("--n-report", type=int, default=20000)
ap.add_argument("--seed", type=int, default=11)
args = ap.parse_args()
model = NnueNet()
ck = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
model.load_state_dict(ck["model"])
model.eval()
rng = np.random.default_rng(args.seed)
n_train = len(np.load(Path(args.train_tensor_dir) / "stm.npy"))
fit_rows = np.sort(rng.choice(n_train, size=min(args.n_fit, n_train), replace=False))
X, y = _batched_accumulator_and_raw(model, args.train_tensor_dir, fit_rows)
mu = X.mean(axis=0)
sigma = X.std(axis=0)
sigma[sigma < 1e-6] = 1e-6 # guard against a dead/always-zero accumulator dim
Xs = (X - mu) / sigma
# Ridge regression, closed form: w = (X^T X + lambda*I)^-1 X^T y,
# bias handled by centering y (mu/sigma already centers X).
y_mean = y.mean()
yc = y - y_mean
dim = Xs.shape[1]
A = Xs.T @ Xs + RIDGE_LAMBDA * np.eye(dim)
w = np.linalg.solve(A, Xs.T @ yc)
bias = y_mean # Xs has zero mean per-column, so w contributes 0 at the mean point
fit_pred = Xs @ w + bias
fit_resid = y - fit_pred
print(f"Fit on {len(fit_rows)} positions ({args.train_tensor_dir}): "
f"mean|resid|={np.mean(np.abs(fit_resid)):.3f}, "
f"p98={np.percentile(np.abs(fit_resid), 98):.3f}")
n_holdout = len(np.load(Path(args.holdout_tensor_dir) / "stm.npy"))
report_rows = np.sort(rng.choice(n_holdout, size=min(args.n_report, n_holdout), replace=False))
Xh, yh = _batched_accumulator_and_raw(model, args.holdout_tensor_dir, report_rows)
Xhs = (Xh - mu) / sigma
holdout_pred = Xhs @ w + bias
holdout_resid = yh - holdout_pred
abs_resid = np.abs(holdout_resid)
print(f"Held out on {len(report_rows)} positions ({args.holdout_tensor_dir}): "
f"mean|resid|={np.mean(abs_resid):.3f}, "
f"p50={np.percentile(abs_resid, 50):.3f}, "
f"p98={np.percentile(abs_resid, 98):.3f}, "
f"max={np.max(abs_resid):.3f}")
print("NOTE: this is an OFFLINE, quiet-filtered-game-position sample -- "
"NOT real search-node traffic. See the 2026-10-20 conversation on "
"why that distinction matters before trusting these residual "
"numbers as a safety bound for a real skip-gate.")
with open(args.outfile, "wb") as f:
f.write(b"NPRB")
f.write(struct.pack("<I", dim))
f.write(mu.astype(np.float32).tobytes())
f.write(sigma.astype(np.float32).tobytes())
f.write(w.astype(np.float32).tobytes())
f.write(struct.pack("<f", float(bias)))
print(f"Wrote {args.outfile}")
if __name__ == "__main__":
main()
|
