train.py

from __future__ import annotations

import argparse
import json
import queue
import random
import sys
import threading
import time
from pathlib import Path

import torch
import torch.nn.functional as F

sys.path.insert(0, str(Path(__file__).resolve().parent))
from model import (  
    make_reduce_cell, make_add_cell, reduce_features, add_features,
)
from data import make_reduce_batch, make_add_batch  

STAGES = [8, 12, 20, 36, 68, 132, 260, 516, 1027]
TOKEN_BUDGET = 40_000
EVAL_EVERY = 500
ADVANCE_THRESHOLD = 0.9995
SOFT_STEPS = 1_500
NOISE = 0.01
MARGIN = 3.0
MARGIN_W = 0.1
AUX_W = 0.3
PLATEAU_STEPS = 12_000
KINDS = ("reduce", "add")


def features_for(kind, b):
    if kind == "reduce":
        return reduce_features(b["x"], b["p"], b["p3"])
    return add_features(b["x"], b["y"], b["g"])


def forward_loss(kind, cell, batch):
    outs = cell.forward_train(features_for(kind, batch))
    tgt = batch["z"]
    loss_bits = F.binary_cross_entropy_with_logits(outs["bits"], tgt)
    sign = 2 * tgt - 1
    loss_margin = F.relu(MARGIN - outs["bits"] * sign).mean()
    if kind == "reduce":
        loss_aux = (
            F.binary_cross_entropy_with_logits(outs["borrow"], batch["borrow"])
            + F.cross_entropy(outs["q"], batch["q"])
        )
    else:
        loss_aux = F.binary_cross_entropy_with_logits(
            outs["carry"], batch["carry"])
    loss = loss_bits + MARGIN_W * loss_margin + AUX_W * loss_aux
    loss = loss + 1e-4 * sum(blk.last_h_l1 for blk in cell.blocks)
    ok_rows = ((outs["bits"] > 0) == (tgt > 0.5)).all(dim=1)
    return loss, {
        "loss": loss.item(), "bits": loss_bits.item(),
        "exact": ok_rows.float().mean().item(),
    }, ok_rows


def bsz_for(n):
    return min(1024, max(8, TOKEN_BUDGET // n))


def fresh_batch(kind, rng, n, bsz, instances=None):
    if kind == "reduce":
        return make_reduce_batch(rng, n, bsz, instances)
    return make_add_batch(rng, n, bsz, instances)


@torch.no_grad()
def evaluate(kind, cell, rng, n, device, n_batches=4):
    cell.eval()
    total, good = 0, 0
    for _ in range(n_batches):
        b = fresh_batch(kind, rng, n, bsz_for(n))
        b.pop("raw")
        b = {k: v.to(device) for k, v in b.items()}
        logits = cell(features_for(kind, b))
        ok = ((logits > 0) == (b["z"] > 0.5)).all(dim=1)
        good += int(ok.sum())
        total += ok.numel()
    cell.train()
    return good / total


def set_mode(cell, mode, noise):
    cell.gate.mode = mode
    for blk in cell.blocks:
        blk.scan_noise = noise


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--resume", default=None)
    ap.add_argument("--steps", type=int, default=600_000)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--out", default="ckpt2")
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--max-stage", type=int, default=None)
    ap.add_argument("--stages", type=int, nargs="*", default=None)
    ap.add_argument("--force-stage", type=int, default=None)
    ap.add_argument("--eval-window", type=int, default=0,
                    help="0 = min over all widths <= stage; K>0 = min over last K stages only")
    args = ap.parse_args()
    global STAGES
    if args.stages:
        STAGES = args.stages
    if args.max_stage is None:
        args.max_stage = len(STAGES) - 1

    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    rng = random.Random(args.seed)
    torch.manual_seed(args.seed)

    cells = {"reduce": make_reduce_cell("soft").to(device),
             "add": make_add_cell("soft").to(device)}
    opts = {k: torch.optim.AdamW(c.parameters(), lr=args.lr, weight_decay=0.0)
            for k, c in cells.items()}

    import copy as _copy
    ema_cells = {k: _copy.deepcopy(c) for k, c in cells.items()}
    EMA_DECAY = 0.999

    @torch.no_grad()
    def ema_update(kind):
        for pe, pr in zip(ema_cells[kind].parameters(),
                          cells[kind].parameters()):
            pe.lerp_(pr, 1.0 - EMA_DECAY)

    stage, step, last_advance_step = 0, 0, 0
    streak = 0
    best_box: dict[int, float] = {}
    out_dir = Path(args.out)
    out_dir.mkdir(exist_ok=True)
    metrics_f = open(out_dir / "metrics.jsonl", "a")

    if args.resume:
        ck = torch.load(args.resume, map_location=device, weights_only=True)
        cells["reduce"].load_state_dict(ck["reduce_state_dict"], strict=False)
        if "add_state_dict" in ck:
            cells["add"].load_state_dict(ck["add_state_dict"], strict=False)
        for k in KINDS:
            if f"{k}_state_dict" in ck:
                sdk = ck.get(f"{k}_ema_state_dict", ck[f"{k}_state_dict"])
                ema_cells[k].load_state_dict(sdk, strict=False)
        for k in KINDS:
            try:
                opts[k].load_state_dict(ck[f"{k}_opt"])
            except Exception:
                print(f"{k}: fresh optimizer", flush=True)
        stage, step = ck["stage"], ck["step"]
        if args.force_stage is not None:
            stage = args.force_stage
        last_advance_step = step
        print(f"resumed at step {step}, stage {stage}", flush=True)

    def set_lr(v):
        for o in opts.values():
            for g in o.param_groups:
                g["lr"] = v

    def get_lr():
        return opts["reduce"].param_groups[0]["lr"]

    set_lr(args.lr)

    def pick_width(stage_now):
        if stage_now == 0 or rng.random() < 0.5:
            n = STAGES[stage_now]
        else:
            n = STAGES[rng.randint(0, stage_now - 1)]
        if rng.random() < 0.3:
            n = max(5, n - rng.randint(0, max(1, n // 8)))
        return n

    q: queue.Queue = queue.Queue(maxsize=4)
    stage_box = {"stage": stage}
    hard_buf: dict[tuple, list] = {}
    buf_lock = threading.Lock()

    def load_mined(kind):
        import json as _json
        path = out_dir / f"mined_{kind}.jsonl"
        if not path.exists():
            return
        try:
            lines = path.read_text().splitlines()[-40000:]
        except Exception:
            return
        per_width: dict[int, list] = {}
        for ln in lines:
            try:
                row = _json.loads(ln)
            except Exception:
                continue
            if kind == "reduce":
                inst = (row["m"], row["x"])
            else:
                inst = (row["x"], row["y"], row["g"])
            per_width.setdefault(row["n"], []).append(inst)
        with buf_lock:
            for n, insts in per_width.items():
                buf = hard_buf.setdefault((kind, n), [])
                buf.extend(insts)
                if len(buf) > 20_000:
                    del buf[: len(buf) - 20_000]

    def producer():
        prng = random.Random(args.seed + 1)
        last_mined = 0.0
        while True:
            if time.time() - last_mined > 120:
                for k in KINDS:
                    load_mined(k)
                last_mined = time.time()
            kind = "reduce" if prng.random() < 0.5 else "add"
            cur_w = STAGES[min(stage_box["stage"], len(STAGES) - 1)] + 8
            with buf_lock:
                widths = [w for (k, w), v in hard_buf.items()
                          if k == kind and len(v) >= 64 and w <= cur_w]
            if widths and prng.random() < 0.30:
                n = prng.choice(widths)
                bsz = bsz_for(n)
                with buf_lock:
                    pool = hard_buf[(kind, n)]
                    replay = [pool[prng.randrange(len(pool))]
                              for _ in range(bsz // 2)]
                b = fresh_batch(kind, prng, n, bsz - len(replay))
                br = fresh_batch(kind, prng, n, len(replay), instances=replay)
                merged = {k: torch.cat([b[k], br[k]])
                          for k in b if k != "raw"}
                merged["raw"] = b["raw"] + br["raw"]
                q.put((kind, n, merged))
            else:
                n = pick_width(stage_box["stage"])
                q.put((kind, n, fresh_batch(kind, prng, n, bsz_for(n))))

    threading.Thread(target=producer, daemon=True).start()

    for c in cells.values():
        c.train()
    t0 = time.time()
    while step < args.steps:
        step += 1
        mode = "soft" if step <= SOFT_STEPS else "ste"
        noise = 0.0 if step <= SOFT_STEPS else NOISE
        for c in cells.values():
            set_mode(c, mode, noise)

        stage_box["stage"] = stage
        kind, n, batch = q.get()
        raw = batch.pop("raw")
        batch = {k: v.to(device) for k, v in batch.items()}
        cell, opt = cells[kind], opts[kind]
        loss, stats, ok_rows = forward_loss(kind, cell, batch)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(cell.parameters(), 1.0)
        opt.step()
        ema_update(kind)

        if step > SOFT_STEPS:
            bad = (~ok_rows).nonzero().flatten().tolist()
            if bad:
                with buf_lock:
                    buf = hard_buf.setdefault((kind, n), [])
                    buf.extend(raw[i] for i in bad)
                    if len(buf) > 20_000:
                        del buf[: len(buf) - 20_000]

        if step % 100 == 0:
            stats.update(step=step, stage=stage, n=n, kind=kind,
                         lr=get_lr(),
                         sps=round(step / (time.time() - t0), 2))
            print(json.dumps(stats), flush=True)
            metrics_f.write(json.dumps(stats) + "\n")
            metrics_f.flush()

        if step % EVAL_EVERY == 0 and step > SOFT_STEPS:
            accs = {}
            for k in KINDS:
                set_mode(ema_cells[k], "hard", 0.0)
                ema_cells[k].eval()
                lo = 0 if args.eval_window == 0 else max(0, stage + 1 - args.eval_window)
                accs[k] = min(
                    evaluate(k, ema_cells[k], rng, STAGES[s], device,
                             n_batches=2 if s < stage else 4)
                    for s in range(lo, stage + 1))
            print(json.dumps({"eval": accs, "stage": stage,
                              "n": STAGES[stage], "step": step}), flush=True)
            metrics_f.write(json.dumps(
                {"eval": accs, "stage": stage, "step": step}) + "\n")
            metrics_f.flush()
            score = min(accs.values())
            if score > best_box.get(stage, 0.0):
                best_box[stage] = score
                torch.save(
                    {"reduce_state_dict": cells["reduce"].state_dict(),
                     "add_state_dict": cells["add"].state_dict(),
                     "reduce_ema_state_dict":
                         ema_cells["reduce"].state_dict(),
                     "add_ema_state_dict": ema_cells["add"].state_dict(),
                     "stage": stage, "step": step, "score": score},
                    out_dir / f"best_stage{stage}.pt",
                )
            ok = score >= ADVANCE_THRESHOLD
            streak = streak + 1 if ok else 0
            if streak >= 2 and stage < min(args.max_stage, len(STAGES) - 1):
                stage += 1
                streak = 0
                last_advance_step = step
                set_lr(args.lr)
                print(f"=== ADVANCE to stage {stage} (N={STAGES[stage]}) ===",
                      flush=True)
            elif step - last_advance_step >= PLATEAU_STEPS and get_lr() > 5.1e-5:
                set_lr(max(5e-5, get_lr() * 0.5))
                last_advance_step = step
                print(f"=== LR DECAY to {get_lr():.2e} (plateau) ===",
                      flush=True)
            torch.save(
                {"reduce_state_dict": cells["reduce"].state_dict(),
                 "add_state_dict": cells["add"].state_dict(),
                 "reduce_ema_state_dict": ema_cells["reduce"].state_dict(),
                 "add_ema_state_dict": ema_cells["add"].state_dict(),
                 "reduce_opt": opts["reduce"].state_dict(),
                 "add_opt": opts["add"].state_dict(),
                 "stage": stage, "step": step},
                out_dir / "latest.pt",
            )
            if step % 5000 == 0:
                import shutil
                shutil.copy(out_dir / "latest.pt",
                            out_dir / f"snap_{step}.pt")
                snaps = sorted(out_dir.glob("snap_*.pt"),
                               key=lambda p: int(p.stem.split("_")[1]))
                for old in snaps[:-6]:
                    old.unlink()

    metrics_f.close()


if __name__ == "__main__":
    main()