Blame for experiments/train_capacity.py

#!/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()