#!/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()