path: root/annotate_pgns/export_nnue_weights.py

#!/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 "NNU4":
    magic:     4 bytes, ASCII "NNU4"
    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.
    calib:     2x float32 (CALIB_A, CALIB_B) -- affine correction so
               the engine's consumed centipawn output tracks
               Stockfish's own real score, not the net's nominal raw
               WDL-sigmoid scale (which drifts run-to-run even with
               an identical K/lambda/loss-exponent recipe -- see
               CalibrateScoreMapping and the 2026-10-20 "keeping this
               stable across retrains" conversation). NnueForward
               applies `(CALIB_A + CALIB_B*fOut) * 100.0f` instead of
               a bare `fOut * 100.0f`.
    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 \
        --holdout-tensor-dir runs/sd10_twic_2M/tensors_holdout
"""
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 CalibrateScoreMapping(model, holdout_tensor_dir, n_sample=20000, seed=7):
    """Fits a simple affine correction (a, b) so that
    a + b*raw tracks Stockfish's own real score (pawns, mover-
    relative) -- not raw's nominal WDL-sigmoid scale. See the
    2026-10-20 "keeping this stable across retrains" conversation:
    a fixed K/lambda/loss-exponent recipe is NOT enough on its own to
    guarantee the engine's consumed cp output stays on a stable scale
    across generations -- training variance (data, convergence) still
    shifts the empirical raw-to-score mapping even with an identical
    recipe. Re-fitting (a, b) every generation and baking the
    correction into the exported format, applied in NnueForward
    itself, decouples engine search-tuning from needing the raw net's
    internal scale to stay perfectly constant forever.

    Fit against the HOLDOUT tensor dir specifically (not train) --
    an honest estimate needs positions the checkpoint wasn't fit on.
    Uses score_probe_val.npy/score_probe_valid.npy (Stockfish's own
    clipped per-position score, and a validity mask excluding exact/
    mate rows -- see build_training_tensors.py), which
    build_training_tensors.py already writes for every tensor dir
    regardless of whether --score-loss-weight was ever used during
    training. Ordinary least squares, free intercept (not forced
    through the origin)."""
    d = Path(holdout_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")
    score_val = np.load(d / "score_probe_val.npy")
    score_valid = np.load(d / "score_probe_valid.npy").astype(bool)

    valid_rows = np.flatnonzero(score_valid)
    if len(valid_rows) == 0:
        sys.exit(f"CalibrateScoreMapping: {d} has no valid score-probe rows "
                  f"(all exact/mate?) -- can't fit a calibration against it")

    rng = np.random.default_rng(seed)
    rows = rng.choice(valid_rows, size=min(n_sample, len(valid_rows)), replace=False)
    rows.sort()

    raws = []
    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))
            raws.append(model(w_idx, w_off, b_idx, b_off, stm_t).numpy())
    raw = np.concatenate(raws).astype(np.float64)
    score = score_val[rows].astype(np.float64)

    # score ~= a + b*raw, ordinary least squares, free intercept.
    b, a = np.polyfit(raw, score, 1)
    resid = score - (a + b * raw)
    print(f"CalibrateScoreMapping: fit on {len(rows)} holdout positions -- "
          f"a={a:.4f}, b={b:.4f}, mean|resid|={np.mean(np.abs(resid)):.3f} pawns, "
          f"resid p98={np.percentile(np.abs(resid), 98):.3f} pawns")
    return float(a), float(b)


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)")
    ap.add_argument("--holdout-tensor-dir", required=True,
                     help="a build_training_tensors.py HOLDOUT output dir "
                          "(not the same as --tensor-dir), used to fit the "
                          "raw-to-real-Stockfish-score affine correction "
                          "(a, b) baked into the exported format -- see "
                          "CalibrateScoreMapping. Must be holdout, not "
                          "train, for an honest (non-overfit) fit.")
    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)
    CALIB_A, CALIB_B = CalibrateScoreMapping(model, args.holdout_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"NNU4")
        f.write(struct.pack("<4I", DIMENSIONS, L1, L2, L3))
        f.write(struct.pack("<f", FT_SCALE))
        f.write(struct.pack("<3f", AS1, AS2, AS3))
        f.write(struct.pack("<3f", WS1, WS2, WS3))
        f.write(struct.pack("<2f", CALIB_A, CALIB_B))
        f.write(ft_weight_q.tobytes())
        f.write(ft_bias_q.tobytes())
        f.write(fc1_w_q.tobytes())
        f.write(sd["fc1.bias"].detach().contiguous().to(torch.float32).numpy().tobytes())
        f.write(fc2_w_q.tobytes())
        f.write(sd["fc2.bias"].detach().contiguous().to(torch.float32).numpy().tobytes())
        f.write(fc3_w_q.tobytes())
        f.write(sd["fc3.bias"].detach().contiguous().to(torch.float32).numpy().tobytes())

    print(f"Wrote {args.outfile} from {args.checkpoint}")
    print(f"  ft_scale={FT_SCALE}")
    print(f"  act_scales (AS1,AS2,AS3)=({AS1:.4f},{AS2:.4f},{AS3:.4f})")
    print(f"  weight_scales (WS1,WS2,WS3)=({WS1:.4f},{WS2:.4f},{WS3:.4f})")


if __name__ == "__main__":
    main()