File size: 5,612 Bytes
2f06744 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | #!/usr/bin/env python3
"""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: # stability: full-size batch
b = synth_batch(1, 8192, 64, CFG, filler=14, generator=g)
else: # targeted: small batch with weak keys
# consistent by construction: include_keys guarantees the weak keys
# are stored exactly once; qpos guarantees every weak key is queried
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): # fresh census probes, 2048 queries each
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: # always keep the best-so-far
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()
|