| |
| """Real-data 259K retrieval at the 0.999 level (test_realtext.py). |
| |
| v3 design (after two measured failures, documented below): |
| |
| * d_model=64 -- the PROVEN architecture. A d_model=128 rebuild failed to |
| transition twice (uniform keys: frozen at h@64 0.3516 for 3.4K steps; the |
| faithful Zipf recipe: frozen at 0.2422 for 3.6K steps -- while the d=64 |
| original transitioned in <400 steps). Width changes re-roll the |
| transition dice; we don't take that gamble. |
| * warm-started from checkpoints/realtext_base.pt (our own 0.875 model) -- no |
| transition needed at all; this run repairs and extends it. |
| * kernel-D spread loss at the TRUE key-token positions (the 0.875 run applied |
| it at `0::2` of the whole sequence -- mostly filler -- so it never acted). |
| * growth to the full 3,548-pair protocol (the old run grew only to 512 pairs, |
| then evaluated at 3,548 -- out of distribution). |
| * the address scale stays at 181 (the checkpoint's final scale; its |
| degradation came from hitting 181 in the last ~150 steps of training, not |
| from the scale itself -- we re-stabilise at 181 for thousands of steps). |
| * long cosine polish at the full pair count, then an alignment phase at the |
| exact eval distribution (filler=70, Q=2,048, depth-stratified, T~=259K). |
| |
| Protocol identical to experiments/eval_realdata.py: wikitext-103 slice, |
| VS=4,094 word vocab, key->bigram-successor pairs (count>=4), real-text filler. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import sys |
| import time |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| ROOT = os.path.join(HERE, "..") |
| sys.path.insert(0, os.path.join(ROOT, "src")) |
|
|
| from mulattice.protocol import RealData, logits_at, stratified_qpos |
| from mulattice.models import MQARModel, MixerCfg |
|
|
| DEVICE = os.environ.get("DEVICE", "cuda") |
| CORPUS = os.path.join(ROOT, "data", "wikitext103_slice.txt") |
| CKPT_WARM = os.environ.get("WARM", os.path.join(ROOT, "checkpoints", "realtext_base.pt")) |
| CKPT_OUT = os.path.join(ROOT, "checkpoints", "realtext.pt") |
| HEARTBEAT = os.path.join(ROOT, "logs", "realtext.heartbeat") |
| HISTORY = os.path.join(ROOT, "results", "realtext_training.json") |
|
|
| D_MODEL = 64 |
| SCALE = 181.0 |
| P_EVAL = 3548 |
| F_EVAL = 70 |
| Q_EVAL = 2048 |
| TOKEN_BUDGET = 16384 |
|
|
| RESTAB = [(64, 400), (128, 400), (256, 800)] |
| GROW = [(512, 1500), (1024, 2000), (2048, 2500), (3548, 2500)] |
| POLISH_STEPS = 1500 |
| ALIGN_STEPS = 60 |
| if os.environ.get("SMOKE"): |
| RESTAB = [(8, 3), (16, 3)] |
| GROW = [(32, 3), (64, 3)] |
| POLISH_STEPS = 3 |
| ALIGN_STEPS = 2 |
|
|
|
|
| def beat(step, txt=""): |
| with open(HEARTBEAT, "w") as f: |
| f.write(f"{step} {txt}\n") |
|
|
|
|
| def spread_loss_real(m, tokens, key_pos, scale, gate_subsample=2048): |
| from mulattice.protocol import _gated_spread |
| mx = m.blocks[0].mixer |
| xc = mx.conv(m.blocks[0].n1(m.emb(tokens))) |
| kaddr = scale * torch.tanh(mx.Wa(xc)) |
| ar = torch.arange(tokens.shape[0], device=tokens.device).unsqueeze(1) |
| addr = kaddr[ar, key_pos] |
| B, N, _ = addr.shape |
| if N > gate_subsample: |
| addr = addr[:, torch.randperm(N, device=addr.device)[:gate_subsample]] |
| return _gated_spread(addr, addr, scale) |
|
|
|
|
| def main(): |
| rd = RealData(CORPUS) |
| V = rd.V |
| print(f"corpus: {rd.n_pairs} kv-pairs (cap {P_EVAL}), filler stream {len(rd.FILLER)}", |
| flush=True) |
| dev = DEVICE |
| torch.manual_seed(0) |
| m = MQARModel(V, kind="mulattice", n_layers=1, |
| cfg=MixerCfg(d_model=D_MODEL, n_heads=2, d_k=D_MODEL // 2, d_v=D_MODEL // 2, |
| addr_scale=SCALE, learn_addr_scale=False, |
| bound_address=True, retention_init=9.0)).to(dev) |
| if os.path.exists(CKPT_WARM): |
| ck = torch.load(CKPT_WARM, map_location="cpu", weights_only=False) |
| m.load_state_dict(ck["model"]) |
| print(f"warm-start from {CKPT_WARM}", flush=True) |
| else: |
| print(f"WARNING: {CKPT_WARM} missing, training from scratch", flush=True) |
| mx = m.blocks[0].mixer |
| mx.scan.crossover = 1e9 |
| opt = torch.optim.AdamW(m.parameters(), lr=1e-3) |
| g = torch.Generator().manual_seed(1) |
| step = 0 |
| t0 = time.time() |
| history = [] |
|
|
| def held_out(pairs, filler, Q=64, B=2, seed=991, zipf=False): |
| m.eval() |
| c = t = 0 |
| with torch.no_grad(): |
| for _ in range(B): |
| b = rd.batch(pairs, filler, Q, seed=seed + _, zipf=zipf) |
| pred = logits_at(m, b["tokens"].to(dev), b["ans_pos"].to(dev)).argmax(-1).cpu() |
| c += int((pred == b["gold"]).sum()); t += pred.numel() |
| m.train() |
| return c / t |
|
|
| def ce_step(pairs, filler, Q, lr_decay=None, zipf=False, use_spread=True): |
| nonlocal step |
| T_est = pairs * (filler + 2) + 2 * Q |
| B = max(1, min(8, TOKEN_BUDGET // T_est)) |
| b = rd.batch(pairs, filler, Q, seed=int(g.initial_seed()) + step, zipf=zipf) |
| tok = b["tokens"].to(dev) |
| logits = logits_at(m, tok, b["ans_pos"].to(dev)) |
| ce = F.cross_entropy(logits.reshape(-1, V), b["gold"].to(dev).reshape(-1)) |
| if use_spread: |
| sl = spread_loss_real(m, tok, b["key_pos"].to(dev), SCALE) |
| loss = ce + sl |
| else: |
| sl = 0.0 |
| loss = ce |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0) |
| opt.step() |
| if lr_decay: |
| lr_decay.step() |
| step += 1 |
| return float(ce.detach()), (float(sl) if use_spread else 0.0), b["L"], B |
|
|
| |
| |
| print(f"RESTAB (pairs,steps): {RESTAB} zipf, f=2, no spread, lr 1e-3", flush=True) |
| for pairs, n in RESTAB: |
| for s in range(n): |
| ce, sl, L, B = ce_step(pairs, 2, 4, zipf=True, use_spread=False) |
| if s % 200 == 0 or s == n - 1: |
| a = held_out(64, 2, zipf=True) |
| print(f" r p{pairs:4d} step {step:6d} ce {ce:.3f} h@64 {a:.4f} " |
| f"{time.time()-t0:5.0f}s", flush=True) |
| beat(step, f"restab p{pairs}") |
| history.append({"step": step, "phase": "restab", "pairs": pairs, |
| "h64": round(a, 4)}) |
|
|
| |
| print(f"GROWTH (pairs,steps): {GROW} zipf/uniform alternation + true-key spread", |
| flush=True) |
| for pairs, n in GROW: |
| good = 0 |
| for s in range(n): |
| ce, sl, L, B = ce_step(pairs, 2, 8, zipf=(step % 2 == 0), use_spread=True) |
| if s % 200 == 0 or s == n - 1: |
| a = held_out(pairs, 2) |
| a64 = held_out(64, 0, zipf=True) |
| print(f" g p{pairs:4d} step {step:6d} ce {ce:.3f} spread {sl:.3f} " |
| f"h@{pairs} {a:.4f} h@64 {a64:.4f} rounds {mx.scan.last_rounds} " |
| f"{time.time()-t0:5.0f}s", flush=True) |
| beat(step, f"grow p{pairs}") |
| history.append({"step": step, "phase": "grow", "pairs": pairs, |
| "acc": round(a, 4), "h64": round(a64, 4)}) |
| with open(HISTORY, "w") as fh: |
| json.dump(history, fh, indent=2) |
| torch.save({"model": m.state_dict(), |
| "proto_cfg": {"V": V, "d_model": D_MODEL, "addr_scale": SCALE, |
| "n_heads": 2, "d_k": D_MODEL // 2, "d_v": D_MODEL // 2}, |
| "cfg": "realtext periodic", "step": step}, |
| os.path.join(ROOT, "checkpoints", "realtext_latest.pt")) |
| good = good + 1 if a >= 0.99 else 0 |
| if good >= 2 and pairs < P_EVAL: |
| print(f" PROMOTE p{pairs} at step {step}", flush=True) |
| break |
|
|
| |
| print(f"POLISH {POLISH_STEPS} steps at P={P_EVAL}, filler cycle [2,2,4,8], cosine 1e-3->1e-4", |
| flush=True) |
| sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=POLISH_STEPS + ALIGN_STEPS, |
| eta_min=1e-4) |
| fcyc = [2, 2, 4, 8] |
| for s in range(POLISH_STEPS): |
| ce, sl, L, B = ce_step(P_EVAL, fcyc[s % len(fcyc)], 16, lr_decay=sched, |
| zipf=False, use_spread=True) |
| if s % 100 == 0 or s == POLISH_STEPS - 1: |
| print(f" polish {s:4d} step {step} ce {ce:.3f} T{L} {time.time()-t0:6.0f}s", flush=True) |
| beat(step, "polish") |
| if (s + 1) % 250 == 0: |
| a = held_out(P_EVAL, 8, Q=128, B=1) |
| print(f" POLISH EVAL acc {a:.4f}", flush=True) |
| history.append({"step": step, "phase": "polish", "acc": round(a, 4)}) |
| with open(HISTORY, "w") as fh: |
| json.dump(history, fh, indent=2) |
| torch.save({"model": m.state_dict(), |
| "proto_cfg": {"V": V, "d_model": D_MODEL, "addr_scale": SCALE, |
| "n_heads": 2, "d_k": D_MODEL // 2, "d_v": D_MODEL // 2}, |
| "cfg": "realtext periodic", "step": step}, |
| os.path.join(ROOT, "checkpoints", "realtext_latest.pt")) |
|
|
| |
| print(f"ALIGN {ALIGN_STEPS} steps at filler={F_EVAL} (Q=64 train-safe; " |
| f"Q={Q_EVAL} only under no_grad)", flush=True) |
| qpos = stratified_qpos(P_EVAL, 64, 4, seed=4242) |
| for s in range(ALIGN_STEPS): |
| b = rd.batch(P_EVAL, F_EVAL, 64, seed=777000 + s, qpos_idx=qpos) |
| tok = b["tokens"].to(dev) |
| logits = logits_at(m, tok, b["ans_pos"].to(dev)) |
| ce = F.cross_entropy(logits.reshape(-1, V), b["gold"].to(dev).reshape(-1)) |
| sl = spread_loss_real(m, tok, b["key_pos"].to(dev), SCALE, gate_subsample=1024) |
| loss = ce + sl |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0) |
| opt.step() |
| sched.step() |
| step += 1 |
| if s % 10 == 0 or s == ALIGN_STEPS - 1: |
| torch.save({"model": m.state_dict(), |
| "proto_cfg": {"V": V, "d_model": D_MODEL, "addr_scale": SCALE, |
| "n_heads": 2, "d_k": D_MODEL // 2, "d_v": D_MODEL // 2}, |
| "cfg": "realtext periodic", "step": step}, |
| os.path.join(ROOT, "checkpoints", "realtext_latest.pt")) |
| print(f" align {s:3d} step {step} ce {ce:.3f} T{b['L']} {time.time()-t0:6.0f}s", |
| flush=True) |
| beat(step, "align") |
|
|
| torch.save({"model": m.state_dict(), |
| "proto_cfg": {"V": V, "d_model": D_MODEL, "addr_scale": SCALE, |
| "n_heads": 2, "d_k": D_MODEL // 2, "d_v": D_MODEL // 2}, |
| "cfg": "realdata_v2 d64 warm256k spread-true-keys growth3548 polish align"}, |
| CKPT_OUT) |
| print(f"saved {CKPT_OUT}", flush=True) |
|
|
| |
| m.eval() |
| for seed in (4242, 4243): |
| qp = stratified_qpos(P_EVAL, Q_EVAL, 4, seed=seed) |
| b = rd.batch(P_EVAL, F_EVAL, Q_EVAL, seed=seed, qpos_idx=qp) |
| with torch.no_grad(): |
| pred = logits_at(m, b["tokens"].to(dev), b["ans_pos"].to(dev)).argmax(-1).cpu() |
| ok = (pred == b["gold"]).squeeze(0) |
| per = Q_EVAL // 4 |
| strata = [round(ok[i * per:(i + 1) * per].float().mean().item(), 4) for i in range(4)] |
| overall = ok.float().mean().item() |
| print(f"SELF-CHECK seed {seed}: overall {overall:.4f} strata {strata} T={b['L']}", |
| flush=True) |
| history.append({"step": step, "final": True, "seed": seed, |
| "overall": round(overall, 4), "strata": strata}) |
| with open(HISTORY, "w") as fh: |
| json.dump(history, fh, indent=2) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
