1#!/usr/bin/env python2"""Stage B (and ablation variants): train the workspace editor + readout with FORCED mixed action budgets on top3of a (by default frozen) backbone. Also usable to pretrain a tiny backbone from scratch (init: scratch) for the4Phase 4 learnability experiments, and to train the baseline adapters (mode: dense / global_recurrence).56Usage: python -m training.train_editor --config configs/experiments/editor_pilot.yaml [--device mps] [--seed 1]7"""8from __future__ import annotations910import argparse, json, math, sys, time11from pathlib import Path1213import numpy as np14import torch1516sys.path.insert(0, str(Path(__file__).resolve().parents[1]))17from prisme.config import PrismeConfig, load_config # noqa: E40218from prisme.model import PrismeModel # noqa: E40219from prisme.losses import lm_loss, token_nll # noqa: E40220from training.common import (Run, TokenBin, cosine_lr, load_checkpoint, load_yaml, make_adamw, peak_memory_bytes, # noqa: E40221 pick_device, save_checkpoint, seed_all, sync)22from data.synthetic_dataset import AnyTokenizer, SyntheticBatches, exact_match, load_items # noqa: E402232425def build_model(exp: dict):26 m = exp["model"]27 if m == "from_checkpoint": # read the exact config saved with the checkpoint (incl. readout_scale)28 m = json.loads((Path(exp["checkpoint"]) / "checkpoint.json").read_text())["config"]29 cfg = load_config(m) if isinstance(m, str) else PrismeConfig.from_dict(m)30 for k, v in exp.get("model_overrides", {}).items():31 setattr(cfg, k, v)32 cfg.validate()33 mode = exp.get("mode", "prisme")34 if mode == "dense":35 cfg.dense_writes = True36 if mode == "global_recurrence":37 from baselines.global_recurrence import GlobalRecurrenceModel38 return cfg, GlobalRecurrenceModel(cfg, **exp.get("baseline_kwargs", {}))39 if mode == "backbone_only":40 cfg.use_workspace = False41 return cfg, PrismeModel(cfg)424344def sample_forced_actions(cfg, B, T, budgets, probs, rng: np.random.Generator, device, per: str = "position"):45 """Actions [B,T,R_max]: budget b ~ probs per position (or per sequence), actions uniformly random with46 singles/pairs balanced; steps >= b are STOP."""47 R = cfg.r_max48 shape = (B, T) if per == "position" else (B, 1)49 b = rng.choice(budgets, size=shape, p=probs)50 b = np.broadcast_to(b, (B, T))51 grp = rng.integers(0, 2, size=(B, T, R))52 single = rng.integers(0, cfg.n_slots, size=(B, T, R))53 pair = rng.integers(cfg.n_slots, cfg.n_actions, size=(B, T, R))54 acts = np.where(grp == 0, single, pair)55 steps = np.arange(R)[None, None, :]56 acts = np.where(steps < b[..., None], acts, cfg.stop_action)57 return torch.from_numpy(acts).to(device), torch.from_numpy(np.ascontiguousarray(b)).to(device)585960@torch.no_grad()61def eval_by_budget(model, cfg, batches_iter_fn, device, budgets, n_batches, seed=0):62 """Mean token loss under random forced actions at each fixed budget (budget applied to all positions)."""63 model.eval()64 res = {}65 for b in budgets:66 rng = np.random.default_rng(seed)67 tot, n = 0.0, 068 for x, y in batches_iter_fn(n_batches):69 acts, _ = sample_forced_actions(cfg, x.shape[0], x.shape[1], [b], [1.0], rng, device)70 logits = model(x, acts if b > 0 else None)["logits"]71 nll = token_nll(logits, y); mask = (y != -100).float()72 tot += float((nll * mask).sum()); n += int(mask.sum())73 res[b] = tot / max(n, 1)74 model.train()75 return res767778def main(argv=None) -> int:79 ap = argparse.ArgumentParser(description=__doc__)80 ap.add_argument("--config", required=True)81 ap.add_argument("--device", default="auto")82 ap.add_argument("--seed", type=int, default=None)83 ap.add_argument("--run-root", default="experiments/runs")84 ap.add_argument("--dry-run", action="store_true")85 args = ap.parse_args(argv)86 exp = load_yaml(args.config)87 seed = args.seed if args.seed is not None else exp.get("seed", 1)88 seed_all(seed)89 device = pick_device(args.device)90 cfg, model = build_model(exp)91 tr = dict(exp["train"])92 if args.dry_run:93 tr.update(steps=3, batch_size=2, eval_every=2, eval_batches=1, ckpt_every=10**9); device = torch.device("cpu")94 model = model.to(device)95 # ---- init ------------------------------------------------------------------------------------------96 init = exp.get("init", {})97 if init.get("backbone_checkpoint"):98 meta = load_checkpoint(init["backbone_checkpoint"], model, strict=False)99 print("loaded backbone from", init["backbone_checkpoint"], "missing:", len(meta["_missing"]), flush=True)100 freeze = init.get("freeze_backbone", True)101 if freeze:102 for p in model.backbone.parameters():103 p.requires_grad_(False)104 trainable = [p for n, p in model.named_parameters() if p.requires_grad and not n.startswith("controller")]105 auto_scale = init.get("auto_readout_scale", False)106 # ---- data ------------------------------------------------------------------------------------------107 data = exp["data"]108 if data["kind"] == "tokens":109 train_bin = TokenBin(data["train_bin"], tr["seq_len"]); val_bin = TokenBin(data["val_bin"], tr["seq_len"])110 train_iter = train_bin.batches(tr["batch_size"], seed, device)111 val_fn = lambda nb: val_bin.sequential_batches(tr["batch_size"], device, nb)112 tok = None; val_sets = {}113 else:114 tok = AnyTokenizer(data.get("tokenizer", "char"))115 assert tok.vocab_size <= cfg.vocab_size, (tok.vocab_size, cfg.vocab_size)116 train_items = load_items(Path(data["splits_dir"]) / "train.jsonl")117 if data.get("max_train_items"):118 k = len(train_items) / data["max_train_items"]119 train_items = [train_items[int(i * k)] for i in range(data["max_train_items"])]120 train_set = SyntheticBatches(train_items, tok, tr["seq_len"])121 train_iter = train_set.random_batches(tr["batch_size"], seed, device, data.get("loss_span", "answer"))122 val_sets = {name: SyntheticBatches(load_items(Path(data["splits_dir"]) / f"{name}.jsonl"), tok, tr["seq_len"])123 for name in data.get("eval_splits", ["dev_iid"])}124 main_val = val_sets[data.get("eval_splits", ["dev_iid"])[0]]125 val_fn = lambda nb: main_val.sequential_batches(tr["batch_size"], device, data.get("loss_span", "answer"), nb * tr["batch_size"]) # stratified subset (every k-th item)126 budgets = exp.get("budgets", [0, 1, 2, 4, 8]); probs = exp.get("budget_probs", [1 / len(budgets)] * len(budgets))127 probs = [p / sum(probs) for p in probs]128 run = Run(exp.get("name", "editor"), {"experiment": exp, "model": cfg.to_dict(), "resolved_train": tr}, root=args.run_root, seed=seed)129 n_params = sum(p.numel() for p in model.parameters()); n_train = sum(p.numel() for p in trainable)130 run.write_manifest(device=str(device), backend=device.type, precision="fp32", n_params=n_params, n_trainable=n_train,131 checkpoint_parent=init.get("backbone_checkpoint"), mode=exp.get("mode", "prisme"), freeze_backbone=freeze)132 if auto_scale and cfg.use_workspace:133 with torch.no_grad():134 x0, _ = next(train_iter)135 cfg.readout_scale = float(model.backbone(x0).float().pow(2).mean().sqrt())136 run.write_manifest(readout_scale=cfg.readout_scale)137 print(f"readout_scale set to backbone residual rms = {cfg.readout_scale:.3f}", flush=True)138 opt = make_adamw(trainable, tr["lr"], tr.get("weight_decay", 0.1))139 rng = np.random.default_rng(seed + 1000)140 steps = tr["steps"]; t0 = time.time(); seen = 0141 model.train()142 for step in range(steps):143 lr = cosine_lr(step, tr.get("warmup", 100), steps, tr["lr"])144 for g in opt.param_groups:145 g["lr"] = lr146 x, y = next(train_iter)147 if cfg.use_workspace:148 acts, b = sample_forced_actions(cfg, x.shape[0], x.shape[1], budgets, probs, rng, device, exp.get("budget_per", "position"))149 out = model(x, acts)150 else:151 out = model(x, None)152 loss = lm_loss(out["logits"], y)153 opt.zero_grad(set_to_none=True)154 loss.backward()155 gn = torch.nn.utils.clip_grad_norm_(trainable, tr.get("grad_clip", 1.0))156 opt.step()157 seen += int((y != -100).sum())158 if step % tr.get("log_every", 20) == 0 or step == steps - 1:159 sync(device); el = time.time() - t0160 run.log(step=step, loss=loss.item(), lr=lr, grad_norm=float(gn), tokens=seen, tok_per_s=seen / max(el, 1e-6))161 print(f"step {step} loss {loss.item():.4f} lr {lr:.2e} gn {float(gn):.2f} ({el:.0f}s)", flush=True)162 if (step + 1) % tr["eval_every"] == 0 or step == steps - 1:163 ev_b = budgets if cfg.use_workspace else [0]164 res = eval_by_budget(model, cfg, val_fn, device, ev_b, tr.get("eval_batches", 10))165 run.log(step=step, **{f"val_loss_b{b}": v for b, v in res.items()})166 print(" val loss by budget:", {b: round(v, 4) for b, v in res.items()}, flush=True)167 if (step + 1) % tr.get("ckpt_every", 10**9) == 0 or step == steps - 1:168 sha = save_checkpoint(run.dir / "checkpoint", model, cfg.to_dict(), {"step": step + 1, "run_id": run.run_id})169 run.log(step=step, checkpoint_sha256=sha)170 # ---- final evaluation on all requested synthetic splits (exact match at several budgets) -----------171 final = {}172 if tok is not None and exp.get("final_exact_match", True):173 for name, vs in val_sets.items():174 final[name] = {}175 for b in (exp.get("em_budgets", [0, 2, 8]) if cfg.use_workspace else [0]):176 rngb = np.random.default_rng(seed + 7)177 def forced(B_, n, b=b, rngb=rngb):178 return sample_forced_actions(cfg, B_, n, [b], [1.0], rngb, device)[0]179 em = exact_match(model, vs, tok, device, max_budget=b, forced_actions_fn=(forced if b > 0 else None),180 batch_size=tr["batch_size"], max_items=exp.get("em_max_items"))181 final[name][b] = {k: v for k, v in em.items() if k != "records"}182 (run.dir / f"em_{name}_b{b}.json").write_text(json.dumps(em["records"]))183 print(f" exact match {name} budget {b}: {em['accuracy']:.4f} (n={em['n']}) by family {em['by_family']}", flush=True)184 sync(device)185 run.finish(train_tokens=seen, peak_memory_bytes=peak_memory_bytes(device), final_exact_match=final,186 metrics_file=str(run.dir / "metrics.jsonl"))187 print("run", run.run_id, "done")188 return 0189190191if __name__ == "__main__":192 raise SystemExit(main())193