Hugging Face's logo

Spaces:
Ex0bit
/
mLattice
Running

App Files Community
mLattice / experiments /train_capacity.py
Ex0bit's picture
Ex0bit
Docker Space: actual demo server, 3 checkpoints, readable dark text
2f06744 verified
Raw History Blame Contribute Delete
14.7 kB
#!/usr/bin/env python3
"""Capacity curriculum: 8,192 pairs at 131K tokens, target >= 0.999 (test_capacity.py).
Recipe (CONTINUE.md step 1):
* V = 16,384 keys + 16,384 values + 2 (VS=16,384 key vocab so 8,192 stored
keys leave half the vocab as non-key filler material);
* addr_scale = 64 -> (2*64+3)^2 = 17,161 cells for 8,192 keys (~2.1 keys/cell
at the old packing, ~0.5 after the spread loss separates them);
* kernel-D gated spread loss on the TRUE key-token addresses at every step;
* curriculum 64 -> 256 -> 512 -> 1024 -> 2048 -> 4096 -> 8192, auto-promotion
when held-out accuracy at the level target is reached twice consecutively;
* warm-started from checkpoints/mqar_base.pt via vocab surgery (from-scratch
training at V>=16K never escaped the chance plateau: logs/vocab8192.log);
* from the 256 level on, 25% of batches carry 50% overwrites (the forgetting
skill trained simultaneously; train_forgetting.py specialises afterwards).
The loss is cross-entropy at the query SEP positions ONLY (the head is applied
to gathered positions), so no (B,T,V) logits are ever materialised.
"""
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, extend_from_checkpoint,
force_fast_scan, logits_at, spread_loss, synth_batch)
DEVICE = os.environ.get("DEVICE", "cuda")
CKPT_WARM = os.path.join(ROOT, "checkpoints", "mqar_base.pt")
CKPT_OUT = os.path.join(ROOT, "checkpoints", "capacity8k.pt")
HEARTBEAT = os.path.join(ROOT, "logs", "capacity8k.heartbeat")
HISTORY = os.path.join(ROOT, "results", "capacity_training.json")
CFG = SynthCfg()
LEVELS = [
# pairs, max_steps, filler choices (dense-dominant), probe filler, quit acc
dict(pairs=64, max_steps=2500, fillers=[0, 0, 2], probe_f=0, quit=0.99),
dict(pairs=256, max_steps=5000, fillers=[0, 0, 2], probe_f=0, quit=0.99),
dict(pairs=512, max_steps=5000, fillers=[0, 0, 4], probe_f=0, quit=0.99),
dict(pairs=1024, max_steps=6000, fillers=[0, 0, 4], probe_f=0, quit=0.99),
dict(pairs=2048, max_steps=7000, fillers=[0, 0, 6], probe_f=0, quit=0.99),
dict(pairs=4096, max_steps=8000, fillers=[0, 0, 6], probe_f=0, quit=0.99),
dict(pairs=8192, max_steps=16000, fillers=[0, 0, 0, 4, 8], probe_f=14, quit=0.999),
]
if os.environ.get("SMOKE"): # fast local validation of the exact code path
LEVELS = [dict(pairs=8, max_steps=11, fillers=[0, 1], probe_f=0, quit=0.99),
dict(pairs=16, max_steps=11, fillers=[0, 1], probe_f=1, quit=0.99)]
OVERWRITE_FRAC = 0.25 # fraction of batches (pairs>=256) that carry 50% overwrites
ALIGN_STEPS = 600 # final f=14 (131K-protocol) alignment steps at the 8192 level
if os.environ.get("SMOKE"):
ALIGN_STEPS = 3
EVAL_EVERY = 500
if os.environ.get("SMOKE"):
EVAL_EVERY = 5
TOKEN_BUDGET = 12288 # halved after the 175GB OOM: fewer cells per step
TOP = LEVELS[-1]["pairs"] # the capacity target level (8,192; small in SMOKE)
START = int(os.environ.get("START_LEVEL", "0"))
if START:
LEVELS = [lv for lv in LEVELS if lv["pairs"] >= START]
TOP = LEVELS[-1]["pairs"]
# SCALE_END env: the top level interpolates the mixer's addr_scale from the
# checkpoint's value to SCALE_END over its first RAMP steps. Measured
# (train_capacity4.log): at addr_scale=64, 8,192 keys pack 0.49 keys/cell and
# the level plateaus at ~0.96 (probe 0.94 at deep decay) while 4,096 keys at
# 0.24 keys/cell reach 0.98 -- density, not the mechanism, is the limiter.
# scale 96 -> (2*96+3)^2 = 38,049 cells -> 0.22 keys/cell at 8,192 keys.
SCALE_END = float(os.environ.get("SCALE_END", "0")) or None
SCALE_RAMP = int(os.environ.get("SCALE_RAMP", "3000"))
history = []
best_acc = {}
_t_save = 0.0
def beat(step, txt=""):
with open(HEARTBEAT, "w") as f:
f.write(f"{step} {time.time():.0f} {txt}\n")
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": "capacity curriculum ret9 spread"}
if extra:
payload.update(extra)
torch.save(payload, path)
def held_out(m, pairs, filler, Q=64, B=2, seed=991):
g = torch.Generator().manual_seed(seed + pairs + filler)
b = synth_batch(B, pairs, Q, CFG, filler=filler, generator=g)
return eval_batch(m, b, DEVICE)["acc"]
def main():
dev = DEVICE
torch.manual_seed(0)
m = build_model(CFG).to(dev)
RESUME = os.environ.get("RESUME", "")
if RESUME and os.path.exists(RESUME):
ck = torch.load(RESUME, map_location="cpu", weights_only=False)
m.load_state_dict(ck["model"])
# restore the TRAINED address scale: addr_scale is config, not weights,
# so without this the model re-ramps 64->SCALE_END and crashes its own
# accuracy on every restart (measured: 0.97 -> 0.34, train_capacity8.log)
saved_scale = ck.get("proto_cfg", {}).get("addr_scale")
if saved_scale:
m.blocks[0].mixer.cfg.addr_scale = float(saved_scale)
print(f"warm-start: resume from {RESUME} (step {ck.get('step')}, "
f"scale {m.blocks[0].mixer.cfg.addr_scale:.0f})", flush=True)
elif os.path.exists(CKPT_WARM):
extend_from_checkpoint(m, CKPT_WARM, NK=CFG.NK, NV=CFG.NV)
print(f"warm-start: vocab surgery from {CKPT_WARM}", flush=True)
else:
print("WARNING: no warm-start checkpoint; training from scratch", flush=True)
force_fast_scan(m)
opt = torch.optim.AdamW(m.parameters(), lr=3e-3)
g = torch.Generator().manual_seed(1)
step = 0
t0 = time.time()
hist_probe_best = 0.0
def train_level(pairs, max_steps, fillers, probe_f, quit_acc, lr_sched=None):
nonlocal step, hist_probe_best
level_t0 = time.time()
good = 0
last_rounds = None
scale0 = float(m.blocks[0].mixer.cfg.addr_scale)
for s in range(max_steps):
if SCALE_END and pairs <= TOP:
frac = min(1.0, s / SCALE_RAMP)
m.blocks[0].mixer.cfg.addr_scale = scale0 + (SCALE_END - scale0) * frac
f = fillers[step % len(fillers)]
# Q<=16 in TRAINING: the fastscan autograd graph holds ~rounds x n_cells
# cell copies; at Q=64 (rounds~67) x 179K cells (stretched grid) that is
# ~175GB total-vm and the host OOM-killer fires (measured, exit 137,
# train_capacity6.log). Large-Q accuracy is measured by no-grad probes.
# Memory-aware Q: at B<=2 (pairs>=2048) cells<=40K so Q=64 costs
# ~20GB (safe) and is NEEDED -- with Q<=16 the CE signal collapses
# and the big levels erode (p2048/p4096 declined to 0.88-0.93,
# train_capacity7.log). Small levels keep B>1 -> Q<=16 for memory.
Q = (64 if step % 4 == 0 else 16) if pairs >= 2048 else (8 if step % 4 else 16)
ow = OVERWRITE_FRAC if (pairs >= 256 and step % 4 == 1) else 0.0
T_est = int(pairs * (1 + ow) * (f + 2)) + 2 * Q
B = max(1, min(32, TOKEN_BUDGET // T_est))
b = synth_batch(B, pairs, Q, CFG, filler=f, overwrite_frac=ow, generator=g)
tok = b["tokens"].to(dev)
ans = b["ans_pos"].to(dev)
logits = logits_at(m, tok, ans)
ce = F.cross_entropy(logits.reshape(-1, CFG.V), b["gold"].to(dev).reshape(-1))
sl = spread_loss(m, tok, b["key_pos"].to(dev), CFG)
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 lr_sched:
lr_sched.step()
step += 1
if s % 100 == 0 or s == max_steps - 1:
sc = m.blocks[0].mixer.scan
print(f" p{pairs:5d} f{f:2d} step {step:6d} ce {float(ce.detach()):.3f} "
f"spread {float(sl.detach()):.3f} rounds {sc.last_rounds} cells {sc.last_n_cells} "
f"B{B} T{b['L']} {time.time()-t0:6.0f}s", flush=True)
beat(step, f"pairs{pairs}")
if (s + 1) % EVAL_EVERY == 0 or s == max_steps - 1:
acc = held_out(m, pairs, fillers[-1], Q=128)
# per-level rolling checkpoints: latest AND best-by-acc, so a
# degraded level tail can never clobber its own best model
save(m, os.path.join(ROOT, "checkpoints", f"capacity8k_p{pairs}_latest.pt"),
{"step": step, "acc": acc})
if acc >= best_acc.get(pairs, 0.0):
best_acc[pairs] = acc
save(m, os.path.join(ROOT, "checkpoints", f"capacity8k_p{pairs}_best.pt"),
{"step": step, "acc": acc})
rec = {"step": step, "pairs": pairs, "acc": round(acc, 4), "t": round(time.time() - t0)}
probe = None
if pairs == 8192 and probe_f > 0:
probe = held_out(m, pairs, probe_f, Q=512, B=1, seed=5)
rec["probe_131k"] = round(probe, 4)
hist_probe_best = max(hist_probe_best, probe)
history.append(rec)
with open(HISTORY, "w") as fh:
json.dump(history, fh, indent=2)
print(f" EVAL p{pairs} f{fillers[-1]} acc {acc:.4f} sc{float(m.blocks[0].mixer.cfg.addr_scale):.0f}"
+ (f" probe@f{probe_f} {probe:.4f}" if probe is not None else "")
+ f" [{time.time()-level_t0:.0f}s in level]", flush=True)
good = good + 1 if acc >= quit_acc else 0
if good >= 2 and pairs < TOP:
print(f" PROMOTE at step {step} (acc {acc:.4f} x2 >= {quit_acc})", flush=True)
break
if pairs == TOP and probe is not None and probe >= 0.999 and acc >= 0.999:
print(f" TARGET REACHED at step {step}: dense {acc:.4f}, 131K probe {probe:.4f}",
flush=True)
return True
return False
print(f"CAPACITY CURRICULUM V={CFG.V} addr_scale={CFG.addr_scale} device={dev}", flush=True)
for i, lv in enumerate(LEVELS):
# PER-LEVEL cosine decay. Measured (train_capacity3.log): with constant
# lr 3e-3 the p512 level oscillated around 0.985 held-out for its whole
# 5,000-step budget with no upward trend (training ce stuck at 0.03-0.05)
# -- the level never squeezes its last errors. The realdata run only
# reached 1.0000 once cosine decay started. Each level now decays
# 3e-3 -> 1e-3 (top level: -> 3e-4 over steps+align).
# Level-scaled entry LR + warmup. A fresh 3e-3 restart at every level
# EXPLODES a converged model at the big levels (B=1-2, ce 0.085 -> 5.4 in
# 100 steps, train_capacity9.log); small levels with B>=5 survive it.
# Entry: 3e-3 (<=1024), 1.5e-3 (2048), 1e-3 (>=4096); 150-step warmup,
# cosine to eta_min over the level (TOP: 3e-4, else entry/3).
import math as _math
lr0 = 3e-3 if lv["pairs"] <= 1024 else (1.5e-3 if lv["pairs"] == 2048 else 1e-3)
eta = 3e-4 if lv["pairs"] == TOP else lr0 / 3.0
total = lv["max_steps"] + (ALIGN_STEPS if lv["pairs"] == TOP else 0)
warm = 150
fmin = eta / lr0
def _lam(ss):
if ss < warm:
return (ss + 1) / warm
prog = min(1.0, (ss - warm) / max(1, total - warm))
return fmin + (1 - fmin) * 0.5 * (1 + _math.cos(_math.pi * prog))
for pg in opt.param_groups:
pg["lr"] = lr0 * (1.0 / warm)
lr_sched = torch.optim.lr_scheduler.LambdaLR(opt, _lam)
done = train_level(lv["pairs"], lv["max_steps"], lv["fillers"], lv["probe_f"], lv["quit"],
lr_sched)
save(m, os.path.join(ROOT, "checkpoints", f"capacity8k_lvl{lv['pairs']}.pt"),
{"step": step})
if done:
break
# final alignment: the 131K filler distribution at memory-safe query counts.
# (Q>=512 training steps OOM: rounds~Q cell-copies in the autograd graph --
# the realdata align died of exactly this. Large-Q accuracy is measured by
# the no-grad probes and the acceptance test, never trained.)
print(f"ALIGN phase: f=14 (131K protocol), {ALIGN_STEPS} steps, Q<=64", flush=True)
for s in range(ALIGN_STEPS):
Q = 64 if s % 3 == 0 else 16
b = synth_batch(1, TOP, Q, CFG, filler=14, overwrite_frac=0.25 if s % 4 == 1 else 0.0,
generator=g)
tok = b["tokens"].to(dev)
logits = logits_at(m, tok, b["ans_pos"].to(dev))
ce = F.cross_entropy(logits.reshape(-1, CFG.V), b["gold"].to(dev).reshape(-1))
sl = spread_loss(m, tok, b["key_pos"].to(dev), CFG, 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()
step += 1
if s % 50 == 0 or s == ALIGN_STEPS - 1:
print(f" align {s:4d} step {step} ce {float(ce.detach()):.3f} T{b['L']} "
f"{time.time()-t0:6.0f}s", flush=True)
beat(step, "align")
acc = held_out(m, TOP, 14, Q=512, B=1, seed=5)
save(m, os.path.join(ROOT, "checkpoints", "capacity8k_latest.pt"), {"step": step})
history.append({"step": step, "align": True, "probe_131k": round(acc, 4)})
with open(HISTORY, "w") as fh:
json.dump(history, fh, indent=2)
print(f" ALIGN probe@131K acc {acc:.4f}", flush=True)
save(m, CKPT_OUT, {"step": step})
print(f"saved {CKPT_OUT} at step {step}, total {time.time()-t0:.0f}s", flush=True)
# final self-check with the test protocol
a1 = held_out(m, TOP, 14, Q=1024, B=1, seed=20261102)
a2 = held_out(m, TOP, 0, Q=1024, B=1, seed=20261103)
print(f"SELF-CHECK 131K protocol {a1:.4f} dense {a2:.4f}", flush=True)
history.append({"final": True, "acc_131k": round(a1, 4), "acc_dense": round(a2, 4)})
with open(HISTORY, "w") as fh:
json.dump(history, fh, indent=2)
if __name__ == "__main__":
main()