#!/usr/bin/env python """Stage B (and ablation variants): train the workspace editor + readout with FORCED mixed action budgets on top of a (by default frozen) backbone. Also usable to pretrain a tiny backbone from scratch (init: scratch) for the Phase 4 learnability experiments, and to train the baseline adapters (mode: dense / global_recurrence). Usage: python -m training.train_editor --config configs/experiments/editor_pilot.yaml [--device mps] [--seed 1] """ from __future__ import annotations import argparse, json, math, sys, time from pathlib import Path import numpy as np import torch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from prisme.config import PrismeConfig, load_config # noqa: E402 from prisme.model import PrismeModel # noqa: E402 from prisme.losses import lm_loss, token_nll # noqa: E402 from training.common import (Run, TokenBin, cosine_lr, load_checkpoint, load_yaml, make_adamw, peak_memory_bytes, # noqa: E402 pick_device, save_checkpoint, seed_all, sync) from data.synthetic_dataset import AnyTokenizer, SyntheticBatches, exact_match, load_items # noqa: E402 def build_model(exp: dict): m = exp["model"] if m == "from_checkpoint": # read the exact config saved with the checkpoint (incl. readout_scale) m = json.loads((Path(exp["checkpoint"]) / "checkpoint.json").read_text())["config"] cfg = load_config(m) if isinstance(m, str) else PrismeConfig.from_dict(m) for k, v in exp.get("model_overrides", {}).items(): setattr(cfg, k, v) cfg.validate() mode = exp.get("mode", "prisme") if mode == "dense": cfg.dense_writes = True if mode == "global_recurrence": from baselines.global_recurrence import GlobalRecurrenceModel return cfg, GlobalRecurrenceModel(cfg, **exp.get("baseline_kwargs", {})) if mode == "backbone_only": cfg.use_workspace = False return cfg, PrismeModel(cfg) def sample_forced_actions(cfg, B, T, budgets, probs, rng: np.random.Generator, device, per: str = "position"): """Actions [B,T,R_max]: budget b ~ probs per position (or per sequence), actions uniformly random with singles/pairs balanced; steps >= b are STOP.""" R = cfg.r_max shape = (B, T) if per == "position" else (B, 1) b = rng.choice(budgets, size=shape, p=probs) b = np.broadcast_to(b, (B, T)) grp = rng.integers(0, 2, size=(B, T, R)) single = rng.integers(0, cfg.n_slots, size=(B, T, R)) pair = rng.integers(cfg.n_slots, cfg.n_actions, size=(B, T, R)) acts = np.where(grp == 0, single, pair) steps = np.arange(R)[None, None, :] acts = np.where(steps < b[..., None], acts, cfg.stop_action) return torch.from_numpy(acts).to(device), torch.from_numpy(np.ascontiguousarray(b)).to(device) @torch.no_grad() def eval_by_budget(model, cfg, batches_iter_fn, device, budgets, n_batches, seed=0): """Mean token loss under random forced actions at each fixed budget (budget applied to all positions).""" model.eval() res = {} for b in budgets: rng = np.random.default_rng(seed) tot, n = 0.0, 0 for x, y in batches_iter_fn(n_batches): acts, _ = sample_forced_actions(cfg, x.shape[0], x.shape[1], [b], [1.0], rng, device) logits = model(x, acts if b > 0 else None)["logits"] nll = token_nll(logits, y); mask = (y != -100).float() tot += float((nll * mask).sum()); n += int(mask.sum()) res[b] = tot / max(n, 1) model.train() return res def main(argv=None) -> int: ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--config", required=True) ap.add_argument("--device", default="auto") ap.add_argument("--seed", type=int, default=None) ap.add_argument("--run-root", default="experiments/runs") ap.add_argument("--dry-run", action="store_true") args = ap.parse_args(argv) exp = load_yaml(args.config) seed = args.seed if args.seed is not None else exp.get("seed", 1) seed_all(seed) device = pick_device(args.device) cfg, model = build_model(exp) tr = dict(exp["train"]) if args.dry_run: tr.update(steps=3, batch_size=2, eval_every=2, eval_batches=1, ckpt_every=10**9); device = torch.device("cpu") model = model.to(device) # ---- init ------------------------------------------------------------------------------------------ init = exp.get("init", {}) if init.get("backbone_checkpoint"): meta = load_checkpoint(init["backbone_checkpoint"], model, strict=False) print("loaded backbone from", init["backbone_checkpoint"], "missing:", len(meta["_missing"]), flush=True) freeze = init.get("freeze_backbone", True) if freeze: for p in model.backbone.parameters(): p.requires_grad_(False) trainable = [p for n, p in model.named_parameters() if p.requires_grad and not n.startswith("controller")] auto_scale = init.get("auto_readout_scale", False) # ---- data ------------------------------------------------------------------------------------------ data = exp["data"] if data["kind"] == "tokens": train_bin = TokenBin(data["train_bin"], tr["seq_len"]); val_bin = TokenBin(data["val_bin"], tr["seq_len"]) train_iter = train_bin.batches(tr["batch_size"], seed, device) val_fn = lambda nb: val_bin.sequential_batches(tr["batch_size"], device, nb) tok = None; val_sets = {} else: tok = AnyTokenizer(data.get("tokenizer", "char")) assert tok.vocab_size <= cfg.vocab_size, (tok.vocab_size, cfg.vocab_size) train_items = load_items(Path(data["splits_dir"]) / "train.jsonl") if data.get("max_train_items"): k = len(train_items) / data["max_train_items"] train_items = [train_items[int(i * k)] for i in range(data["max_train_items"])] train_set = SyntheticBatches(train_items, tok, tr["seq_len"]) train_iter = train_set.random_batches(tr["batch_size"], seed, device, data.get("loss_span", "answer")) val_sets = {name: SyntheticBatches(load_items(Path(data["splits_dir"]) / f"{name}.jsonl"), tok, tr["seq_len"]) for name in data.get("eval_splits", ["dev_iid"])} main_val = val_sets[data.get("eval_splits", ["dev_iid"])[0]] 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) budgets = exp.get("budgets", [0, 1, 2, 4, 8]); probs = exp.get("budget_probs", [1 / len(budgets)] * len(budgets)) probs = [p / sum(probs) for p in probs] run = Run(exp.get("name", "editor"), {"experiment": exp, "model": cfg.to_dict(), "resolved_train": tr}, root=args.run_root, seed=seed) n_params = sum(p.numel() for p in model.parameters()); n_train = sum(p.numel() for p in trainable) run.write_manifest(device=str(device), backend=device.type, precision="fp32", n_params=n_params, n_trainable=n_train, checkpoint_parent=init.get("backbone_checkpoint"), mode=exp.get("mode", "prisme"), freeze_backbone=freeze) if auto_scale and cfg.use_workspace: with torch.no_grad(): x0, _ = next(train_iter) cfg.readout_scale = float(model.backbone(x0).float().pow(2).mean().sqrt()) run.write_manifest(readout_scale=cfg.readout_scale) print(f"readout_scale set to backbone residual rms = {cfg.readout_scale:.3f}", flush=True) opt = make_adamw(trainable, tr["lr"], tr.get("weight_decay", 0.1)) rng = np.random.default_rng(seed + 1000) steps = tr["steps"]; t0 = time.time(); seen = 0 model.train() for step in range(steps): lr = cosine_lr(step, tr.get("warmup", 100), steps, tr["lr"]) for g in opt.param_groups: g["lr"] = lr x, y = next(train_iter) if cfg.use_workspace: acts, b = sample_forced_actions(cfg, x.shape[0], x.shape[1], budgets, probs, rng, device, exp.get("budget_per", "position")) out = model(x, acts) else: out = model(x, None) loss = lm_loss(out["logits"], y) opt.zero_grad(set_to_none=True) loss.backward() gn = torch.nn.utils.clip_grad_norm_(trainable, tr.get("grad_clip", 1.0)) opt.step() seen += int((y != -100).sum()) if step % tr.get("log_every", 20) == 0 or step == steps - 1: sync(device); el = time.time() - t0 run.log(step=step, loss=loss.item(), lr=lr, grad_norm=float(gn), tokens=seen, tok_per_s=seen / max(el, 1e-6)) print(f"step {step} loss {loss.item():.4f} lr {lr:.2e} gn {float(gn):.2f} ({el:.0f}s)", flush=True) if (step + 1) % tr["eval_every"] == 0 or step == steps - 1: ev_b = budgets if cfg.use_workspace else [0] res = eval_by_budget(model, cfg, val_fn, device, ev_b, tr.get("eval_batches", 10)) run.log(step=step, **{f"val_loss_b{b}": v for b, v in res.items()}) print(" val loss by budget:", {b: round(v, 4) for b, v in res.items()}, flush=True) if (step + 1) % tr.get("ckpt_every", 10**9) == 0 or step == steps - 1: sha = save_checkpoint(run.dir / "checkpoint", model, cfg.to_dict(), {"step": step + 1, "run_id": run.run_id}) run.log(step=step, checkpoint_sha256=sha) # ---- final evaluation on all requested synthetic splits (exact match at several budgets) ----------- final = {} if tok is not None and exp.get("final_exact_match", True): for name, vs in val_sets.items(): final[name] = {} for b in (exp.get("em_budgets", [0, 2, 8]) if cfg.use_workspace else [0]): rngb = np.random.default_rng(seed + 7) def forced(B_, n, b=b, rngb=rngb): return sample_forced_actions(cfg, B_, n, [b], [1.0], rngb, device)[0] em = exact_match(model, vs, tok, device, max_budget=b, forced_actions_fn=(forced if b > 0 else None), batch_size=tr["batch_size"], max_items=exp.get("em_max_items")) final[name][b] = {k: v for k, v in em.items() if k != "records"} (run.dir / f"em_{name}_b{b}.json").write_text(json.dumps(em["records"])) print(f" exact match {name} budget {b}: {em['accuracy']:.4f} (n={em['n']}) by family {em['by_family']}", flush=True) sync(device) run.finish(train_tokens=seen, peak_memory_bytes=peak_memory_bytes(device), final_exact_match=final, metrics_file=str(run.dir / "metrics.jsonl")) print("run", run.run_id, "done") return 0 if __name__ == "__main__": raise SystemExit(main())