1#!/usr/bin/env python2"""Exact parameter accounting for PRISME configurations (meta-device instantiation).34Reports unique parameters per component and reconciles against the reference counts of the WP31 brief.5Usage: python scripts/count_params.py --all-configs --device meta [--json results/tables/param_counts.json]6"""7from __future__ import annotations89import argparse, json, sys10from pathlib import Path1112import torch1314sys.path.insert(0, str(Path(__file__).resolve().parents[1]))15from prisme.config import PrismeConfig, load_config # noqa: E40216from prisme.model import PrismeModel # noqa: E4021718REFERENCE = { # from the WP31 brief (reference design, pre-implementation)19 "prisme_103m": {"backbone": 99_502_848, "workspace": 3_088_260, "total": 102_591_108},20 "prisme_459m": {"backbone": 446_744_064, "workspace": 12_336_900, "total": 459_080_964},21 "prisme_1180m": {"backbone": 1_149_339_648, "workspace": 25_779_844, "total": 1_175_119_492},22}232425def closed_form_backbone(c: PrismeConfig) -> int:26 d, f, V, L = c.d_model, c.ffn_width, c.vocab_size, c.n_layers27 kv = c.kv_dim28 per_layer = 2 * d * d + 2 * d * kv + 3 * d * f + 2 * d29 return V * d + L * per_layer + d303132def closed_form_workspace(c: PrismeConfig) -> dict[str, int]:33 d, w, fw, M = c.d_model, c.w, c.ffn_w, c.n_slots34 hidden = c.controller_hidden or fw35 parts = {36 "bank_norm": d,37 "bank_projections (W_K, W_V, P)": 3 * d * w,38 "slot_embed e": M * w,39 "editor.block (shared x Lw)": 6 * w * w + 3 * w * fw + 3 * w,40 "editor.ctx_proj": 3 * w + w,41 "editor.gate": 2 * w + 1,42 "readout (q,k,U)": 2 * w * w + w * d,43 "controller": 2 * w * w + (2 * w + w) + w + 3 * w * hidden + 2 * w + (4 * w + 4),44 }45 return parts464748def reference_closed_form(c: PrismeConfig) -> int:49 """Algebraic reconstruction of the brief's workspace/controller count from its three published values:50 4*d*w + M*w + 17*w^2 + 6*w*f_w + 15.5*w + 4 (fits all three configurations exactly)."""51 d, w, fw, M = c.d_model, c.w, c.ffn_w, c.n_slots52 return int(4 * d * w + M * w + 17 * w * w + 6 * w * fw + 15.5 * w + 4)535455def count(model: torch.nn.Module) -> dict[str, int]:56 seen, out = set(), {}57 for n, p in model.named_parameters():58 if id(p) in seen:59 continue60 seen.add(id(p))61 top = n.split(".")[0]62 out[top] = out.get(top, 0) + p.numel()63 return out646566def main() -> int:67 ap = argparse.ArgumentParser(description=__doc__)68 ap.add_argument("--all-configs", action="store_true")69 ap.add_argument("--config", nargs="*", default=[])70 ap.add_argument("--device", default="meta")71 ap.add_argument("--json", default="results/tables/param_counts.json")72 args = ap.parse_args()73 paths = sorted(Path("configs").glob("prisme_*.yaml")) if args.all_configs else [Path(p) for p in args.config]74 report = {}75 for p in paths:76 cfg = load_config(p)77 with torch.device(args.device):78 model = PrismeModel(cfg)79 by_top = count(model)80 backbone = by_top.get("backbone", 0)81 workspace = sum(v for k, v in by_top.items() if k != "backbone")82 total = backbone + workspace83 cf_b = closed_form_backbone(cfg)84 cf_w = closed_form_workspace(cfg)85 ref = REFERENCE.get(cfg.name)86 rec = {87 "config": str(p), "name": cfg.name, "config_sha256": cfg.sha256(),88 "by_module": by_top, "backbone": backbone, "workspace_controller": workspace, "total": total,89 "closed_form_backbone": cf_b, "closed_form_workspace_parts": cf_w, "closed_form_workspace": sum(cf_w.values()),90 "reference_brief": ref,91 "reference_workspace_closed_form_reconstruction": reference_closed_form(cfg),92 "backbone_matches_reference": (ref is not None and backbone == ref["backbone"]),93 "workspace_delta_vs_reference": (workspace - ref["workspace"]) if ref else None,94 }95 assert backbone == cf_b, (backbone, cf_b)96 assert workspace == sum(cf_w.values()), (workspace, sum(cf_w.values()))97 report[cfg.name] = rec98 print(f"{cfg.name}: backbone={backbone:,} workspace+controller={workspace:,} total={total:,} "99 f"| ref backbone={ref['backbone']:,} ({'OK' if rec['backbone_matches_reference'] else 'MISMATCH'}), "100 f"ref workspace={ref['workspace']:,} (delta {rec['workspace_delta_vs_reference']:+,})")101 Path(args.json).parent.mkdir(parents=True, exist_ok=True)102 Path(args.json).write_text(json.dumps(report, indent=2))103 return 0104105106if __name__ == "__main__":107 raise SystemExit(main())108