| |
| """ |
| 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) |
| |
| 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") |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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() |
|
|
| |
| 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) |
| |
| 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() |
|
|
