#!/usr/bin/env python3 """Exports a trained nnue_model.py checkpoint into a flat binary file the C engine (src/nnueeval.c's NnueLoadWeights) can read directly -- no parsing, just a fixed-order sequence of arrays behind a small header. File format (all little-endian), version "NNU3": magic: 4 bytes, ASCII "NNU3" dims: 4x uint32 -- NNUE_DIMENSIONS, L1, L2, L3 (sanity-checked against the C side's own compile-time constants on load, so a mismatched build can't silently misread the file) ft_scale: 1x float32 -- fixed-point scale for ft.weight/ft_bias (int16). NnueForward divides the int32 accumulator sum by this to get the float "fc1 input" value. act_scale: 3x float32 (AS1, AS2, AS3) -- fixed-point scales used to quantize each stage's post-ReLU activation to uint8 before feeding the next layer's int8 dot product: AS1 quantizes fc1's input (the ReLU'd, concatenated accumulator halves) AS2 quantizes fc1's ReLU'd output (fc2's input) AS3 quantizes fc2's ReLU'd output (fc3's input) Derived from real sampled activations (see CalibrateActivationScales) with headroom, NOT from a theoretical max -- see docs/NNUE.md's "Quantization plan" / the int8 float-vs-quantized correctness check. w_scale: 3x float32 (WS1, WS2, WS3) -- fixed-point scales for fc1.weight/fc2.weight/fc3.weight (int8), one per layer, each derived from that layer's own max|.| (exact, not sampled -- weights are fixed, known precisely from the checkpoint). then, in this exact order: ft.weight [NNUE_DIMENSIONS x L1] int16 (quantized @ ft_scale) ft_bias [L1] int16 (quantized @ ft_scale) fc1.weight [L2 x (2*L1)] int8 (quantized @ WS1) fc1.bias [L2] float32 (unquantized -- added post-dequant, tiny array, not worth it) fc2.weight [L3 x L2] int8 (quantized @ WS2) fc2.bias [L3] float32 fc3.weight [1 x L3] int8 (quantized @ WS3) fc3.bias [1] float32 Both the embedding table (int16) and the FC head (int8) are quantized now -- see docs/NNUE.md's "Quantization plan". Measured accuracy cost of the FC-head int8 step, 5000 real holdout positions vs the exact float reference: mean |delta| 1.55cp, median 0.99cp, max 20.91cp, zero clipping -- an order of magnitude below the sd10-vs-sd16 search-depth noise (mean 19cp) the project already tolerates in its training labels. See the conversation/commit history for the simulation that produced those numbers; CalibrateActivationScales below reproduces the scale derivation deterministically from a real checkpoint + its training tensors. Usage: python3 export_nnue_weights.py \ runs/sd10_twic_2M/nnue_ckpt/ckpt_epoch019.pt \ runs/sd10_twic_2M/nnue_ckpt/weights.nnue \ --tensor-dir runs/sd10_twic_2M/tensors_train """ import argparse import struct import sys from pathlib import Path import numpy as np import torch from nnue_feature_spec import DIMENSIONS, PS_NB from nnue_model import L1, L2, L3, NnueNet # Fixed-point scale for ft.weight/ft_bias. Checked against # ckpt_epoch019.pt (ft.weight max|.|=5.01, ft_bias max|.|=0.52): at # this scale, 5.01*4096 ~= 20502, comfortably under int16's 32767 # ceiling (~1.6x headroom for a larger Phase B net before this needs # re-deriving). See docs/NNUE.md's "Quantization plan". FT_SCALE = 4096.0 # Headroom multiplier applied to the largest value seen in a real # sample before deriving each activation scale -- protects against # rare tail values in the full corpus that a finite sample might not # have hit. 1.3x chosen as a round, conservative margin; revisit if a # larger/differently-distributed Phase B net ever reports clipping. ACT_HEADROOM = 1.3 def CalibrateActivationScales(model, tensor_dir, n_sample=50000, seed=42): """Runs n_sample real positions from tensor_dir (one of build_training_tensors.py's output directories) through the float model, capturing each stage's real post-ReLU activation distribution via forward hooks, and returns (AS1, AS2, AS3) -- uint8 fixed-point scales with ACT_HEADROOM margin over the observed max. Must run against the SAME checkpoint being exported -- these scales are not portable across checkpoints.""" 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") n = len(stm) rng = np.random.default_rng(seed) rows = rng.choice(n, size=min(n_sample, n), replace=False) rows.sort() captured = {} h1 = model.fc1.register_forward_hook( lambda m, i, o: captured.setdefault("fc1_in", []).append(i[0].detach())) h2 = model.fc2.register_forward_hook( lambda m, i, o: captured.setdefault("fc2_in", []).append(i[0].detach())) h3 = model.fc3.register_forward_hook( lambda m, i, o: captured.setdefault("fc3_in", []).append(i[0].detach())) with torch.no_grad(): for i in range(0, len(rows), 5000): chunk = rows[i:i + 5000] 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)) model(w_idx, w_off, b_idx, b_off, stm_t) h1.remove(); h2.remove(); h3.remove() scales = [] for key in ("fc1_in", "fc2_in", "fc3_in"): t = torch.cat(captured[key]).float() m = t.max().item() scales.append(255.0 / (m * ACT_HEADROOM)) return tuple(scales) def QuantizeFt(arr): """Rounds + clamps a float32 array to int16 at FT_SCALE. Returns (int16 ndarray, number of values that clipped).""" scaled = np.round(arr.astype(np.float64) * FT_SCALE) clipped = int(np.count_nonzero((scaled < -32768) | (scaled > 32767))) scaled = np.clip(scaled, -32768, 32767) return scaled.astype(np.int16), clipped def QuantizeWeightS8(arr, scale): """Rounds + clamps a float32 array to int8 at the given scale. Returns (int8 ndarray, number of values that clipped).""" scaled = np.round(arr.astype(np.float64) * scale) clipped = int(np.count_nonzero((scaled < -128) | (scaled > 127))) scaled = np.clip(scaled, -128, 127) return scaled.astype(np.int8), clipped def CoalesceFactorization(sd, tensor_dir, n_check=4096, seed=123): """If the checkpoint has a trained ft_p factorization table (see nnue_model.py's 2026-10-18 note), folds it into ft.weight in place (idx % PS_NB is the P-factor's row for real feature idx) and returns a modified copy of sd with ft_p dropped -- the coalesced ft alone must reproduce the factored forward pass, so ft_p carries zero inference-time cost. No-op (returns sd unchanged) if ft_p isn't present (an older, unfactored checkpoint). Required correctness check, per NNUE.md: a bad coalesce silently produces a WORSE net (real rows learn less because the factor absorbed shared signal) with no sign of it in the training loss curve. Verified here by comparing the pre-coalesce (factored, two-table) forward pass against the post-coalesce (plain, ft-only) forward pass on n_check real positions sampled from tensor_dir -- must agree to ~1e-5 or this refuses to export.""" if "ft_p.weight" not in sd: return sd ft_p_weight = sd["ft_p.weight"] if tuple(ft_p_weight.shape) != (PS_NB, L1): sys.exit(f"ft_p.weight: expected shape {(PS_NB, L1)}, checkpoint " f"has {tuple(ft_p_weight.shape)} -- wrong model version?") idx_mod = torch.arange(DIMENSIONS) % PS_NB coalesced_ft_weight = sd["ft.weight"] + ft_p_weight[idx_mod] sd_factored = dict(sd) model_factored = NnueNet() model_factored.load_state_dict(sd_factored) model_factored.eval() sd_plain = dict(sd) sd_plain["ft.weight"] = coalesced_ft_weight sd_plain["ft_p.weight"] = torch.zeros_like(ft_p_weight) model_plain = NnueNet() model_plain.load_state_dict(sd_plain) model_plain.eval() 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") rng = np.random.default_rng(seed) rows = rng.choice(len(stm), size=min(n_check, len(stm)), replace=False) rows.sort() w_idx_list, w_off_list, b_idx_list, b_off_list = [], [], [], [] off_w, off_b = 0, 0 for r in rows: 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[rows].astype(np.int64)) with torch.no_grad(): out_factored = model_factored(w_idx, w_off, b_idx, b_off, stm_t) out_plain = model_plain(w_idx, w_off, b_idx, b_off, stm_t) max_delta = (out_factored - out_plain).abs().max().item() if max_delta > 1e-4: sys.exit(f"Coalesce check FAILED: max |factored - coalesced| = " f"{max_delta:.6g} over {n_check} real positions (expected " f"~1e-5) -- refusing to export a bad coalesce. Check " f"ft_p's shape/init and the idx % PS_NB formula.") print(f"Coalesce check OK: max |factored - coalesced| = {max_delta:.3g} " f"over {n_check} real positions") sd_out = dict(sd) sd_out["ft.weight"] = coalesced_ft_weight del sd_out["ft_p.weight"] return sd_out def main(): ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("checkpoint") ap.add_argument("outfile") ap.add_argument("--tensor-dir", required=True, help="a build_training_tensors.py output dir, used " "to calibrate the FC head's activation scales " "against real positions (e.g. tensors_train)") args = ap.parse_args() ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=True) sd = ckpt["model"] sd = CoalesceFactorization(sd, args.tensor_dir) expected_shapes = { "ft.weight": (DIMENSIONS, L1), "ft_bias": (L1,), "fc1.weight": (L2, 2 * L1), "fc1.bias": (L2,), "fc2.weight": (L3, L2), "fc2.bias": (L3,), "fc3.weight": (1, L3), "fc3.bias": (1,), } for name, shape in expected_shapes.items(): if tuple(sd[name].shape) != shape: sys.exit(f"{name}: expected shape {shape}, checkpoint has " f"{tuple(sd[name].shape)} -- wrong model version?") model = NnueNet() # strict=False: CoalesceFactorization deliberately drops ft_p.weight # (folded into ft.weight above) when the checkpoint had it, but # NnueNet always declares the ft_p parameter regardless -- missing_keys # is expected to be exactly ["ft_p.weight"] in that case, or empty for # an older unfactored checkpoint CoalesceFactorization no-op'd on. missing, unexpected = model.load_state_dict(sd, strict=False) bad_missing = [k for k in missing if k != "ft_p.weight"] if bad_missing or unexpected: sys.exit(f"load_state_dict mismatch beyond the expected ft_p.weight " f"drop -- missing={bad_missing}, unexpected={unexpected}") model.eval() AS1, AS2, AS3 = CalibrateActivationScales(model, args.tensor_dir) ft_weight_f32 = sd["ft.weight"].detach().contiguous().to(torch.float32).numpy() ft_bias_f32 = sd["ft_bias"].detach().contiguous().to(torch.float32).numpy() fc1_w_f32 = sd["fc1.weight"].detach().contiguous().to(torch.float32).numpy() fc2_w_f32 = sd["fc2.weight"].detach().contiguous().to(torch.float32).numpy() fc3_w_f32 = sd["fc3.weight"].detach().contiguous().to(torch.float32).numpy() ft_weight_q, clipped_ftw = QuantizeFt(ft_weight_f32) ft_bias_q, clipped_ftb = QuantizeFt(ft_bias_f32) WS1 = 127.0 / np.abs(fc1_w_f32).max() WS2 = 127.0 / np.abs(fc2_w_f32).max() WS3 = 127.0 / np.abs(fc3_w_f32).max() fc1_w_q, clipped_w1 = QuantizeWeightS8(fc1_w_f32, WS1) fc2_w_q, clipped_w2 = QuantizeWeightS8(fc2_w_f32, WS2) fc3_w_q, clipped_w3 = QuantizeWeightS8(fc3_w_f32, WS3) total_clipped = (clipped_ftw + clipped_ftb + clipped_w1 + clipped_w2 + clipped_w3) if total_clipped: sys.exit(f"Quantize: {total_clipped} value(s) clipped -- " f"re-check headroom/scales before shipping this export " f"(see docs/NNUE.md)") with open(args.outfile, "wb") as f: f.write(b"NNU3") f.write(struct.pack("<4I", DIMENSIONS, L1, L2, L3)) f.write(struct.pack("