| |
| """Weak-key targeted polish (final 0.2 points). |
| |
| Census (results/weak_keys.json, 2 seeds x 4,096 queries): ALL residual errors |
| at the 0.998 plateau come from ~16 distinct key ids out of 16,384 -- 0.1% of |
| the vocabulary. The full-batch polish cannot move them (0.4% exposure per |
| step) and uniform training wobbles the optimum instead. |
| |
| This script: |
| * alternates small 512-pair batches whose stored keys and query set FORCE the |
| weak keys (every batch stores AND queries all of them; DFT 1/p weighting |
| on top) with full 8,192-pair stability batches; |
| * probes on FRESH census seeds (never trained on) every 100 steps, f=14, |
| Q=2,048, and saves when two consecutive probes show <= 1 error total. |
| """ |
| 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 (SynthCfg, build_model, eval_batch, force_fast_scan, logits_at, |
| spread_loss, synth_batch) |
|
|
| DEVICE = os.environ.get("DEVICE", "cuda") |
| CKPT_IN = os.environ.get("RESUME", os.path.join(ROOT, "checkpoints", "capacity8k.pt")) |
| CKPT_OUT = os.environ.get("OUT", os.path.join(ROOT, "checkpoints", "capacity8k.pt")) |
| MAX_STEPS = int(os.environ.get("STEPS", "500")) |
| LR = float(os.environ.get("LR", "5e-5")) |
| CFG = SynthCfg() |
|
|
|
|
| def save(m, path, extra=None): |
| live_scale = float(m.blocks[0].mixer.cfg.addr_scale) |
| payload = {"model": m.state_dict(), "proto_cfg": { |
| "NK": CFG.NK, "NV": CFG.NV, "d_model": CFG.d_model, "n_heads": CFG.n_heads, |
| "d_k": CFG.d_k, "d_v": CFG.d_v, "addr_scale": live_scale, |
| "retention_init": CFG.retention_init}, "cfg": "weak-key targeted polish"} |
| if extra: |
| payload.update(extra) |
| torch.save(payload, path) |
|
|
|
|
| def weighted_ce(logits, gold): |
| p_gold = torch.softmax(logits, -1).gather(1, gold.unsqueeze(1)).squeeze(1) |
| w = (1.0 / p_gold.clamp_min(1e-4)).clamp(1.0, 10.0) |
| ce_vec = F.cross_entropy(logits, gold, reduction="none") |
| return (w * ce_vec).sum() / w.sum() |
|
|
|
|
| def main(): |
| dev = DEVICE |
| weak = json.load(open(os.path.join(ROOT, "results", "weak_keys.json"))) |
| weak_t = torch.tensor(sorted(weak)) |
| print(f"weak-key polish: {len(weak_t)} keys, lr {LR}, max {MAX_STEPS}", flush=True) |
| ck = torch.load(CKPT_IN, map_location="cpu", weights_only=False) |
| cfg = SynthCfg() |
| if ck.get("proto_cfg", {}).get("addr_scale"): |
| cfg.addr_scale = float(ck["proto_cfg"]["addr_scale"]) |
| m = build_model(cfg) |
| m.load_state_dict(ck["model"]) |
| m = m.to(dev).train() |
| force_fast_scan(m) |
| opt = torch.optim.AdamW(m.parameters(), lr=LR) |
| g = torch.Generator().manual_seed(21) |
| t0 = time.time() |
| good = 0 |
| probe_seed = 301 |
| hist = [] |
| for s in range(MAX_STEPS): |
| if s % 3 == 2: |
| b = synth_batch(1, 8192, 64, CFG, filler=14, generator=g) |
| else: |
| |
| |
| nw = len(weak_t) |
| qpos = torch.cat([torch.arange(nw), |
| nw + torch.randperm(512 - nw, generator=g)[: 128 - nw] |
| ]).unsqueeze(0) |
| b = synth_batch(1, 512, 128, CFG, filler=14, generator=g, qpos=qpos, |
| kw={"include_keys": weak_t}) |
| tok = b["tokens"].to(dev) |
| logits = logits_at(m, tok, b["ans_pos"].to(dev)) |
| gold = b["gold"].to(dev).reshape(-1) |
| ce = weighted_ce(logits.reshape(-1, CFG.V), gold) |
| sl = spread_loss(m, tok, b["key_pos"].to(dev), CFG, gate_subsample=512) |
| loss = ce + sl |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0) |
| opt.step() |
| if s % 50 == 0: |
| print(f" wk {s:4d} ce {float(ce.detach()):.4f} T{b['L']} {time.time()-t0:6.0f}s", |
| flush=True) |
| if (s + 1) % 100 == 0 or s == MAX_STEPS - 1: |
| errs = 0 |
| for _ in range(2): |
| gg = torch.Generator().manual_seed(probe_seed) |
| probe_seed += 1 |
| bb = synth_batch(1, 8192, 1024, CFG, filler=14, generator=gg) |
| m.eval() |
| with torch.no_grad(): |
| pred = logits_at(m, bb["tokens"].to(dev), bb["ans_pos"].to(dev)).argmax(-1).cpu() |
| m.train() |
| errs += int((pred != bb["gold"]).sum()) |
| hist.append({"s": s, "errs_2048": errs}) |
| with open(os.path.join(ROOT, "results", "polish_weakkeys.json"), "w") as fh: |
| json.dump(hist, fh, indent=2) |
| print(f" CENSUS probe (2x1024 fresh queries): {errs} errors " |
| f"[{time.time()-t0:.0f}s]", flush=True) |
| good = good + 1 if errs <= 1 else 0 |
| if errs <= 3: |
| save(m, CKPT_OUT, {"step": s, "census_errs": errs}) |
| if good >= 2: |
| print(f"TARGET: two consecutive censuses <= 1 error at step {s}", flush=True) |
| save(m, CKPT_OUT, {"step": s, "census_errs": errs}) |
| return |
| print("done", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
