| """ |
| healthcheck.py -- checkpoint health report for Quazimoto-LM. |
| |
| Auto-finds the checkpoint (same rule as train.py --resume: the --out/--ckpt file |
| if present, else the newest *.pt in the folder), then reports: |
| 1. load + global tensor health (NaN/Inf, dtypes, param count, step) |
| 2. per-(grouped)-layer weight stats with dead/exploded flags |
| 3. family-trait layer status -- each trait's gates (tanh) so you can see which |
| blocks have actually woken up, plus their key weight stats |
| 4. current perplexity on a sample of text (the local --data file, or a built-in |
| fallback string), with the random-vocab baseline for reference. |
| |
| Usage: |
| python healthcheck.py # auto-find in ./chkpt |
| python healthcheck.py --ckpt path/to.pt |
| python healthcheck.py --data corpus.txt # measure PPL on your own text |
| """ |
|
|
| import argparse |
| import math |
| import os |
| import re |
| import sys |
| from collections import Counter |
|
|
| import torch |
|
|
| from model import QuazimotoLM, QuazimotoConfig |
| from family import (HRMRefinementBlock, JEPAPredictorBlock, MoESwiGLU, MTPHead, |
| EngramRing, PhaseAttentionRing, RingControllerBank, RingSpecialists) |
|
|
| PKG_DIR = os.path.dirname(os.path.abspath(__file__)) |
|
|
|
|
| |
| |
| |
| def find_ckpt(path): |
| """Prefer an exact file; else newest *.pt in its folder (or ./chkpt).""" |
| if path and os.path.isfile(path): |
| return path |
| folder = os.path.dirname(path) if path else os.path.join(PKG_DIR, "chkpt") |
| folder = folder or "." |
| if not os.path.isdir(folder): |
| return None |
| pts = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".pt")] |
| return max(pts, key=os.path.getmtime) if pts else None |
|
|
|
|
| def load_tokenizer(tok_dir): |
| sys.path.insert(0, tok_dir) |
| from spike_tokenizer import SpikeTokenizer |
| return SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json")) |
|
|
|
|
| def section(title): |
| print("\n" + "=" * 72 + f"\n{title}\n" + "=" * 72) |
|
|
|
|
| def fmt_gate(t): |
| return "[" + ", ".join(f"{x:+.3f}" for x in t) + "]" |
|
|
|
|
| |
| |
| |
| def report_load(ckpt, path): |
| section("1. LOAD + GLOBAL TENSOR HEALTH") |
| print(f"checkpoint : {path}") |
| print(f"keys : {list(ckpt.keys())}") |
| print(f"step : {ckpt.get('step')} | saved vocab_size: {ckpt.get('vocab_size')}") |
| print(f"optimizer : {'present (resumable)' if 'optim' in ckpt else 'absent'}") |
| sd = ckpt["model"] |
| nan = inf = total = 0 |
| dtypes = Counter() |
| for v in sd.values(): |
| if not torch.is_tensor(v): |
| continue |
| total += v.numel(); dtypes[str(v.dtype)] += 1 |
| nan += torch.isnan(v).sum().item(); inf += torch.isinf(v).sum().item() |
| print(f"tensors : {len(sd)} | params: {total/1e6:.1f}M | dtypes: {dict(dtypes)}") |
| flag = "OK" if (nan == 0 and inf == 0) else "*** PROBLEM ***" |
| print(f"NaNs : {nan} | Infs: {inf} -> {flag}") |
| return sd |
|
|
|
|
| |
| |
| |
| def report_layers(sd): |
| section("2. PER-LAYER WEIGHT HEALTH (grouped over repeated layers)") |
| buckets = {} |
| for k, v in sd.items(): |
| if not torch.is_tensor(v) or v.dtype != torch.float32: |
| continue |
| base = re.sub(r"\.\d+\.", ".N.", k) |
| buckets.setdefault(base, []).append(v) |
| print(f"{'param group':46s} {'std':>8} {'absmax':>9} {'mean':>9} flag") |
| for base, vs in buckets.items(): |
| if "ring_bank" in base or "specialists" in base: |
| continue |
| s = torch.cat([x.flatten() for x in vs]).float() |
| if s.numel() < 2: |
| continue |
| flag = "" |
| if s.std() < 1e-4 and s.abs().max() < 1e-4: flag = "DEAD?" |
| if s.abs().max() > 50: flag = "EXPLODED?" |
| print(f"{base:46s} {s.std():8.4f} {s.abs().max():9.3f} {s.mean():+9.4f} {flag}") |
|
|
|
|
| |
| |
| |
| def _wstat(t): |
| t = t.float() |
| return f"std {t.std():.4f} absmax {t.abs().max():.3f}" if t.numel() > 1 else f"val {t.item():+.4f}" |
|
|
|
|
| def report_traits(model, cfg): |
| section("3. FAMILY-TRAIT LAYER STATUS (tanh-gate ~0 => still a no-op)") |
| any_trait = False |
|
|
| |
| for li, layer in enumerate(model.layers): |
| q = layer.quaz |
| tag = f"layer{li}.quaz" |
| print(f"\n[{tag}] oscillator block gate tanh = {torch.tanh(q.gate).item():+.4f} " |
| f"| block_gain {_wstat(q.block_gain)} | alpha {_wstat(q.alpha)}") |
| if hasattr(q, "phase_seed_gate"): |
| print(f" fractal phase-seed gate : {torch.tanh(q.phase_seed_gate).item():+.4f} " |
| f"(0 => seed unused)") |
| if getattr(q, "use_rings", False): |
| ar = [torch.tanh(r.gate).item() for r in q.attn_rings] |
| er = [torch.tanh(r.scale).item() for r in q.engram_rings] |
| print(f" phase-attn ring gates : {fmt_gate(ar)}") |
| print(f" engram ring scales : {fmt_gate(er)}") |
| if getattr(q, "use_ring_specialists", False): |
| sc = [torch.tanh(b.scale).item() for b in q.specialists] |
| act = [b.active.mean().item() for b in q.specialists] |
| filled = [float(b.store_out.abs().sum() > 0) for b in q.specialists] |
| print(f" specialist out-gates : {fmt_gate(sc)}") |
| print(f" specialist mean 'active': {fmt_gate(act)}") |
| print(f" specialist store filled : {fmt_gate(filled)} (0 until written at inference)") |
| if li == 0 and (getattr(q, 'use_rings', False) or getattr(q, 'use_ring_specialists', False)): |
| any_trait = True |
|
|
| |
| if model.hrm is not None: |
| any_trait = True |
| g = torch.tanh(model.hrm.gate) |
| print(f"\n[HRM] gate tanh = {fmt_gate(g.tolist())} " |
| f"({'OPEN' if g.abs().mean() > 0.01 else 'no-op'}) | " |
| f"z0 {_wstat(model.hrm.z0)} | down {_wstat(model.hrm.down.weight)}") |
| if model.moe is not None: |
| any_trait = True |
| print(f"\n[MoE] router {_wstat(model.moe.router.weight)} | " |
| f"experts={model.moe.n_routed} shared={model.moe.n_shared} top_k={model.moe.top_k}") |
| if model.mtp_heads is not None: |
| any_trait = True |
| norms = [h.proj.weight.norm().item() for h in model.mtp_heads] |
| print(f"\n[MTP] {len(model.mtp_heads)} draft heads | proj-weight norms " |
| f"{fmt_gate(norms)} (0 => head still identity / untrained)") |
| if model.jepa is not None: |
| any_trait = True |
| g = torch.tanh(model.jepa.gate) |
| print(f"\n[JEPA] gate tanh = {fmt_gate(g.tolist())} (train-only aux; never runs in generate)") |
| if model.ring_bank is not None: |
| any_trait = True |
| decs = [c.dec.weight.norm().item() for c in model.ring_bank.controllers] |
| print(f"\n[Controllers] {len(model.ring_bank.controllers)} ring controllers | " |
| f"decoder-weight norms {fmt_gate(decs)} (0 => no-op self-organizer)") |
| if getattr(cfg, "use_fractal_phase_seed", False): |
| any_trait = True |
| gates = [torch.tanh(layer.quaz.phase_seed_gate).item() for layer in model.layers] |
| mag = sum(abs(g) for g in gates) / max(len(gates), 1) |
| tbl = os.path.join(PKG_DIR, "fractal_phase.pt") |
| if os.path.exists(tbl): |
| d = torch.load(tbl, map_location="cpu", weights_only=False) |
| ok = tuple(d.get("phases").shape) == (cfg.vocab_size, cfg.n_osc) |
| tinfo = f"{d.get('mode','?')} table {'MATCHES' if ok else '** SHAPE MISMATCH **'} " \ |
| f"{tuple(d.get('phases').shape)}" |
| else: |
| tinfo = "no fractal_phase.pt found -> model uses FLAT Halton fallback" |
| print(f"\n[FractalSeed] per-layer phase-seed gate tanh {fmt_gate(gates)}") |
| print(f" status: {'OPEN' if mag > 0.01 else 'no-op (gates ~0)'} | {tinfo}") |
|
|
| if not any_trait: |
| print("\n(no opt-in family traits enabled in this checkpoint)") |
|
|
|
|
| |
| |
| |
| @torch.no_grad() |
| def report_ppl(model, cfg, tok, data_path, device): |
| section("4. PERPLEXITY") |
| if data_path and os.path.exists(data_path): |
| with open(data_path, "r", encoding="utf-8", errors="replace") as f: |
| text = f.read()[:200000] |
| src = data_path |
| else: |
| text = ("the quazimoto oscillator learns to synchronize phases across rings. " |
| "coupled clocks fall into step when coupling is strong enough. ") * 50 |
| src = "built-in fallback string" |
| ids = tok.encode(text, add_special_tokens=False) |
| ids = ids[: cfg.block_size * 8 + 1] |
| if len(ids) < 2: |
| print("not enough tokens to score."); return |
| x = torch.tensor([ids], device=device) |
| losses = [] |
| B = cfg.block_size |
| for s in range(0, x.size(1) - 1, B): |
| xb = x[:, s:s + B]; yb = x[:, s + 1:s + 1 + B] |
| n = min(xb.size(1), yb.size(1)) |
| if n < 1: break |
| _, loss, _ = model(xb[:, :n], yb[:, :n]) |
| losses.append(loss.item()) |
| ce = sum(losses) / len(losses) |
| print(f"text source : {src} ({len(ids)} tokens, {len(losses)} windows)") |
| print(f"CE loss : {ce:.4f}") |
| print(f"bits/token : {ce/math.log(2):.4f}") |
| print(f"perplexity : {math.exp(ce):.1f} (random-vocab baseline: {cfg.vocab_size})") |
|
|
|
|
| |
| def main(): |
| p = argparse.ArgumentParser(description="Quazimoto-LM checkpoint health check") |
| p.add_argument("--ckpt", default="", help="checkpoint path (default: auto-find newest in ./chkpt)") |
| p.add_argument("--data", default="", help="text file to measure PPL on (default: built-in string)") |
| p.add_argument("--tok_dir", default=PKG_DIR) |
| p.add_argument("--device", default="cpu") |
| args = p.parse_args() |
|
|
| path = find_ckpt(args.ckpt) |
| if path is None: |
| print("No checkpoint found. Pass --ckpt or train one first.") |
| return |
|
|
| ckpt = torch.load(path, map_location=args.device, weights_only=False) |
| sd = report_load(ckpt, path) |
| report_layers(sd) |
|
|
| cfg = QuazimotoConfig(**ckpt["family_config"]) |
| model = QuazimotoLM(cfg) |
| miss, unexp = model.load_state_dict(sd, strict=False) |
| model.to(args.device).eval() |
| if miss or unexp: |
| print(f"\n[load] missing {len(miss)} / unexpected {len(unexp)} keys " |
| f"(arch drift vs this code; e.g. {(miss or unexp)[:2]})") |
|
|
| tok = load_tokenizer(args.tok_dir) |
| report_traits(model, cfg) |
| report_ppl(model, cfg, tok, args.data, args.device) |
| print() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
