experiments/exp_precision.py

#!/usr/bin/env python3
"""
Adaptive numerical precision (Phase 11).

Distinguishes the four things the brief says must not be conflated:

  A. static quantisation            - one level everywhere
  B. operation-dependent precision  - per-site levels (attn/local/mlp/mem)
  C. step-dependent precision       - cheap early steps, precise late steps
  D. position-dependent precision   - precision chosen per position from a
                                      difficulty signal (policy entropy)

Every reduced-precision path is validated NUMERICALLY against the FP32
reference, and WALL-CLOCK is measured -- fake-quant does not make CPU maths
faster, and saying so is part of the result.

    python experiments/exp_precision.py --checkpoint runs/x/latest.pt
"""
from __future__ import annotations

import argparse
import json
import os
import random
import sys
import time

import chess
import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from tinychess.batching import encode_batch
from tinychess.checkpoint import load_model_only
from tinychess.quant import LEVELS, QuantPolicy, quantise_model_weights


def positions(n=64, seed=0):
    rng = random.Random(seed)
    out = []
    while len(out) < n:
        b = chess.Board()
        for _ in range(rng.randint(0, 50)):
            ms = list(b.legal_moves)
            if not ms:
                break
            b.push(rng.choice(ms))
        if not b.is_game_over():
            out.append(b)
    return out


@torch.no_grad()
def run(model, pb, steps, qp=None, reps=3):
    t0 = time.time()
    for _ in range(reps):
        out = model(**pb.model_args(), steps=steps, qp=qp)
    dt = (time.time() - t0) / reps
    return out, dt


def compare(ref, got, mask):
    """Agreement between a reference and a reduced-precision run."""
    lr = torch.log_softmax(ref["logits"], -1)
    lg = torch.log_softmax(got["logits"], -1)
    pr, pg = lr.exp(), lg.exp()
    top_ref = ref["logits"].argmax(-1)
    top_got = got["logits"].argmax(-1)
    valid = torch.isfinite(lr) & torch.isfinite(lg)
    kl = float((pr[valid] * (lr[valid] - lg[valid])).sum() / ref["logits"].shape[0])
    return {
        "top1_agreement": float((top_ref == top_got).float().mean()),
        "kl_from_reference": kl,
        "max_abs_logit_err": float((ref["logits"][valid] - got["logits"][valid]).abs().max()),
        "value_mae": float((ref["value"] - got["value"]).abs().mean()),
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", required=True)
    ap.add_argument("--steps", type=int, default=6)
    ap.add_argument("--n", type=int, default=64)
    ap.add_argument("--out", default="runs/precision.json")
    a = ap.parse_args()

    model, cfg = load_model_only(a.checkpoint)
    # forward() clamps n to cfg.max_steps; lift the cap so --steps is honoured.
    trained_max_steps = int(cfg.max_steps)
    if a.steps > trained_max_steps and hasattr(model, "core"):
        import torch.nn as nn
        old = model.core.step_emb.data
        new = torch.zeros(a.steps + 1, old.shape[1])
        new[:old.shape[0]] = old
        new[old.shape[0]:] = old[-1]
        model.core.step_emb = nn.Parameter(new)
        model.cfg.max_steps = a.steps
        model.core.cfg.max_steps = a.steps
        print(f"[lifted] step cap {trained_max_steps} -> {a.steps}")
    boards = positions(a.n)
    pb = encode_batch(boards)
    res = {"checkpoint": a.checkpoint, "steps": a.steps, "n_positions": len(boards)}

    ref, t_ref = run(model, pb, a.steps, None)
    res["reference"] = {"seconds": round(t_ref, 4), "dtype": "fp32"}
    print(f"FP32 reference: {t_ref*1000:.1f} ms/forward")

    # ---------------------------------------------- A. static activation quant
    print("\n[A] static activation precision (all sites the same)")
    static = {}
    for lvl in ("bf16", "fp16", "int8", "int4", "int3", "int2"):
        qp = QuantPolicy(default=lvl)
        got, dt = run(model, pb, a.steps, qp)
        m = compare(ref, got, pb.cand_mask)
        m["seconds"] = round(dt, 4)
        m["speed_vs_fp32"] = round(t_ref / dt, 3)
        static[lvl] = m
        print(f"  {lvl:5s} top1={m['top1_agreement']:.3f} KL={m['kl_from_reference']:.5f} "
              f"maxerr={m['max_abs_logit_err']:.4f} {dt*1000:.1f}ms")
    res["static_activation"] = static

    # ---------------------------------------------- weight quantisation
    print("\n[A2] static WEIGHT quantisation (destructive, restored after)")
    wq = {}
    sd = {k: v.clone() for k, v in model.state_dict().items()}
    for bits, name in ((8, "int8"), (4, "int4"), (3, "int3"), (2, "int2")):
        err = quantise_model_weights(model, bits)
        got, dt = run(model, pb, a.steps, None)
        m = compare(ref, got, pb.cand_mask)
        m["mean_weight_rel_error"] = round(sum(err.values()) / max(len(err), 1), 5)
        wq[name] = m
        print(f"  {name:5s} top1={m['top1_agreement']:.3f} KL={m['kl_from_reference']:.5f} "
              f"w_err={m['mean_weight_rel_error']:.4f}")
        model.load_state_dict(sd)
    res["static_weights"] = wq

    # ---------------------------------------------- B. operation-dependent
    print("\n[B] operation-dependent precision")
    op_res = {}
    for name, sites in (
        ("cheap_local_precise_attn", {"local": "int3", "mlp": "int4", "attn": "fp32"}),
        ("cheap_attn_precise_mlp", {"attn": "int3", "local": "int4", "mlp": "fp32"}),
        ("all_int4_except_mem", {"attn": "int4", "local": "int4", "mlp": "int4"}),
    ):
        qp = QuantPolicy(default="fp32", sites=sites)
        got, dt = run(model, pb, a.steps, qp)
        m = compare(ref, got, pb.cand_mask)
        m["sites"] = sites
        op_res[name] = m
        print(f"  {name:26s} top1={m['top1_agreement']:.3f} KL={m['kl_from_reference']:.5f}")
    res["operation_dependent"] = op_res

    # ---------------------------------------------- C. step-dependent
    print("\n[C] step-dependent precision (cheap early, precise late)")

    class StepPolicy(QuantPolicy):
        """Precision as a function of the recurrent step."""
        def __init__(self, early="int3", late="fp32", switch=3):
            super().__init__(default=early)
            self.early, self.late, self.switch, self.step = early, late, switch, 0

        def level_for(self, site):
            return self.late if self.step >= self.switch else self.early

    sp = StepPolicy()
    orig = model.core.forward

    def wrapped(x, mem, step, qp=None, collect=None):
        if isinstance(qp, StepPolicy):
            qp.step = step
        return orig(x, mem, step, qp, collect)

    model.core.forward = wrapped
    got, dt = run(model, pb, a.steps, sp)
    m = compare(ref, got, pb.cand_mask)
    res["step_dependent"] = {"early_int3_late_fp32_switch3": m}
    print(f"  int3 (steps<3) -> fp32: top1={m['top1_agreement']:.3f} "
          f"KL={m['kl_from_reference']:.5f}")
    model.core.forward = orig

    # ---------------------------------------------- D. position-dependent
    print("\n[D] position-dependent precision (difficulty = policy entropy)")
    lp = torch.log_softmax(ref["logits"], -1)
    p = lp.exp()
    ent = -(p * torch.where(torch.isfinite(lp), lp, torch.zeros_like(lp))).sum(-1)
    thr = ent.median()
    easy = (ent <= thr).nonzero().flatten()
    hard = (ent > thr).nonzero().flatten()

    # easy positions get int3, hard positions get fp32
    sub_easy = encode_batch([boards[i] for i in easy.tolist()])
    sub_hard = encode_batch([boards[i] for i in hard.tolist()])
    ge, te = run(model, sub_easy, a.steps, QuantPolicy(default="int3"))
    gh, th = run(model, sub_hard, a.steps, None)
    re_, _ = run(model, sub_easy, a.steps, None)
    rh_, _ = run(model, sub_hard, a.steps, None)
    m_easy = compare(re_, ge, sub_easy.cand_mask)
    m_hard = compare(rh_, gh, sub_hard.cand_mask)
    # counterfactual: what if we had used int3 on the HARD ones instead?
    gh3, _ = run(model, sub_hard, a.steps, QuantPolicy(default="int3"))
    m_hard_cheap = compare(rh_, gh3, sub_hard.cand_mask)
    res["position_dependent"] = {
        "n_easy": len(easy), "n_hard": len(hard),
        "easy_at_int3": m_easy, "hard_at_fp32": m_hard,
        "hard_at_int3_counterfactual": m_hard_cheap,
        "int3_hurts_hard_more": m_hard_cheap["top1_agreement"] < m_easy["top1_agreement"],
    }
    print(f"  easy@int3 top1={m_easy['top1_agreement']:.3f} | "
          f"hard@int3 top1={m_hard_cheap['top1_agreement']:.3f} "
          f"-> cheap precision hurts hard positions more: "
          f"{res['position_dependent']['int3_hurts_hard_more']}")

    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    json.dump(res, open(a.out, "w"), indent=2, default=str)
    print(f"\n[saved] {a.out}")


if __name__ == "__main__":
    main()