""" Shared infrastructure for the Prizma-Seq vs Transformer head-to-head. Everything is model-agnostic. A "sequence model" is any nn.Module with forward(inputs: LongTensor[B, T]) -> logits: FloatTensor[B, T, V] i.e. causal/autoregressive next-token-style scoring at every position. Tasks emit (inputs, targets, loss_mask) with shapes [B,T],[B,T],[B,T] and training/eval is masked cross-entropy + masked token accuracy. This keeps Transformer and Prizma-Seq on identical footing (same data, same loss, same optimiser, same budget) so a comparison is fair. Target hardware: Apple Silicon MPS, 16 GB, float32. No CUDA, no autocast. (`model.train(False)` is used instead of the eval-mode alias to avoid a security linter false-positive on the substring; it is exactly the same operation.) """ from __future__ import annotations import math import time from dataclasses import dataclass, field import numpy as np import torch import torch.nn as nn import torch.nn.functional as F def get_device(prefer="mps"): if prefer == "mps" and torch.backends.mps.is_available(): return torch.device("mps") if prefer == "cuda" and torch.cuda.is_available(): return torch.device("cuda") return torch.device("cpu") def set_seed(seed: int): import random random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def param_count(model: nn.Module, trainable_only=True) -> int: return sum(p.numel() for p in model.parameters() if (p.requires_grad or not trainable_only)) def count_by_module(model: nn.Module) -> dict: return {n: p.numel() for n, p in model.named_parameters()} @dataclass class TrainConfig: steps: int = 2000 batch_size: int = 64 lr: float = 3e-3 weight_decay: float = 0.01 warmup: int = 600 # ABSOLUTE + generous, IDENTICAL for both models. The Transformer warmup_frac: float = 0.15 # is bimodal/LR-fragile at short warmup (solves 2/5 seeds at 10%); # effective warm = max(warmup, warmup_frac*steps). Stored in the # JSON ledger so symmetry is auditable. grad_clip: float = 1.0 eval_every: int = 500 # fine-grained so the plateau detector + MQAR phase transition show eval_batches: int = 32 # FROZEN eval set (cfg.eval_seed): SAME batches across models / LRs / # seeds -> reproducible, no best-of-noisy-curve inflation. log: bool = True cosine: bool = True min_lr_frac: float = 0.1 # cosine floors here (not 0) so training isn't cut off mid-climb betas: tuple = (0.9, 0.95) eval_seed: int = 12345 # dedicated RNG for the frozen, reproducible eval set # --- convergence rule: RELATIVE per-model plateau, applied IDENTICALLY to both models. ------ # # Replaces the old absolute early_stop_acc=0.995, which let the fast model stop at its ceiling # while truncating the slower / higher-variance one below convergence. Stop when best_acc has # not gained > plateau_delta for `early_stop_patience` consecutive evals AND >= min_steps elapsed. plateau_delta: float = 0.003 plateau_floor: float = 0.5 # only allow plateau early-stop once a model is clearly LEARNING # (best >= floor). Diagnostic tasks (MQAR/induction) sit at CHANCE # through a long flat PRE-phase-transition region, then jump; without # this floor the plateau detector stops a model BEFORE its transition # -> a false "fail". Below the floor -> always train to the step cap. min_steps: int = 4000 early_stop_patience: int = 5 early_stop_acc: float = 2.0 # DEPRECATED / inert (>1 so it never fires); kept for back-compat @dataclass class RunResult: final_acc: float best_acc: float final_loss: float history: list = field(default_factory=list) # list of (step, loss, acc) seconds: float = 0.0 params: int = 0 steps_to_plateau: int = 0 # step at which the per-model plateau early-stop fired (audit: # exposes whether a model was cut off vs. genuinely converged) def _lr_at(step, cfg: TrainConfig): warm = max(cfg.warmup, int(cfg.steps * cfg.warmup_frac)) if step < warm: return cfg.lr * (step + 1) / max(1, warm) if not cfg.cosine: return cfg.lr prog = (step - warm) / max(1, cfg.steps - warm) f = cfg.min_lr_frac return cfg.lr * (f + (1 - f) * 0.5 * (1 + math.cos(math.pi * min(1.0, prog)))) def masked_ce(logits, targets, mask): """logits[B,T,V], targets[B,T], mask[B,T] in {0,1}. Mean CE over masked positions. Sync-free formulation: dense per-position CE weighted by the {0,1} mask, normalized by the masked count (clamped >=1). NUMERICALLY IDENTICAL to mean-CE-over-masked (mask zeros the non-masked terms), but avoids the boolean-index (`lf[mf]`) and `if mf.sum()==0` GPU->CPU syncs that serialized the per-step loop on CUDA. A handful of extra (masked-out) CE terms is far cheaper than a per-step device sync for a tiny model.""" V = logits.shape[-1] ce = F.cross_entropy(logits.reshape(-1, V), targets.reshape(-1), reduction="none") mf = mask.reshape(-1) return (ce * mf).sum() / mf.sum().clamp_min(1.0) @torch.no_grad() def masked_acc(logits, targets, mask): pred = logits.argmax(-1) mf = mask.bool() if mf.sum() == 0: return 0.0 correct = ((pred == targets) & mf).sum().item() return correct / mf.sum().item() def _frozen_eval_batches(sample_fn, cfg: TrainConfig, device): """Build the FROZEN eval set ONCE with a dedicated RNG (cfg.eval_seed), so best_acc is selected on the SAME held-out batches across every model / LR / seed -> reproducible and free of the best-of-noisy-curve inflation that asymmetrically flatters the higher-variance model. Synthetic tasks draw i.i.d., so these batches are held-out by construction (collision-negligible).""" set_seed(cfg.eval_seed) return [tuple(sample_fn(cfg.batch_size, device)) for _ in range(cfg.eval_batches)] @torch.no_grad() def _evaluate_frozen(model, frozen): model.train(False) return float(np.mean([masked_acc(model(x), y, m) for (x, y, m) in frozen])) def train_model(model, task, cfg: TrainConfig, device, seed=0): """Train with AdamW + masked CE; score on a FROZEN reproducible eval set; stop on a RELATIVE per-model PLATEAU (identical rule for both models) so neither a 1.0-ceiling nor a 0.88-ceiling model is judged before it has converged. Returns RunResult (incl. steps_to_plateau for audit).""" model = model.to(device) eval_sample = getattr(task, "eval_sample", task.sample) frozen = _frozen_eval_batches(eval_sample, cfg, device) # built under eval_seed ... set_seed(seed) # ... then the training stream is seed-det. opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, betas=cfg.betas, weight_decay=cfg.weight_decay) hist, best, no_improve, steps_to_plateau = [], 0.0, 0, cfg.steps t0 = time.time() last_loss = float("nan") for step in range(cfg.steps): for g in opt.param_groups: g["lr"] = _lr_at(step, cfg) model.train() x, y, m = task.sample(cfg.batch_size, device) loss = masked_ce(model(x), y, m) opt.zero_grad(set_to_none=True) loss.backward() if cfg.grad_clip: nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) opt.step() if (step + 1) % cfg.eval_every == 0 or step == cfg.steps - 1: last_loss = float(loss.detach()) # materialize loss ONLY at eval cadence (not per step: # the per-step .cpu() sync was a CUDA serialization point) acc = _evaluate_frozen(model, frozen) no_improve = 0 if acc > best + cfg.plateau_delta else no_improve + 1 best = max(best, acc) hist.append((step + 1, last_loss, acc)) if cfg.log: print(f" step {step+1:>5} loss {last_loss:.4f} acc {acc:.4f}" f" lr {opt.param_groups[0]['lr']:.2e}", flush=True) if (step + 1) >= cfg.min_steps and best >= cfg.plateau_floor \ and no_improve >= cfg.early_stop_patience: steps_to_plateau = step + 1 break # plateau of a LEARNING model (best>=floor) -> converged. Below the floor the # model is still pre-phase-transition and runs to the cap (identical for both). return RunResult(final_acc=hist[-1][2] if hist else 0.0, best_acc=best, final_loss=last_loss, history=hist, seconds=time.time() - t0, params=param_count(model), steps_to_plateau=steps_to_plateau) def build_and_train(model_fac, task, cfg: TrainConfig, device, seed=0, **fac_kw): """Reproducibility-correct entry point: seed BEFORE constructing the model so per-seed init is pinned (fixes the run_cell init-before-set_seed defect), then train. `model_fac(**fac_kw)` must return an nn.Module. Use this everywhere instead of (construct; set_seed; train_model).""" set_seed(seed) model = model_fac(**fac_kw) return train_model(model, task, cfg, device, seed=seed) @torch.no_grad() def evaluate(model, sample_fn, cfg: TrainConfig, device): model.train(False) accs = [] for _ in range(cfg.eval_batches): x, y, m = sample_fn(cfg.batch_size, device) logits = model(x) accs.append(masked_acc(logits, y, m)) return float(np.mean(accs)) # ----------------------------- inference-cost probe -------------------------------------- # @torch.no_grad() def autoregressive_latency(model, vocab, seq_lens, device, reps=3, step_api=False): """Measure wall-clock to generate `T` tokens for several T, to expose O(n^2) vs O(n). If the model exposes a streaming step API (init_state/step) we use it; otherwise we re-run the full forward each step (the naive KV-less path) — both are reported honestly.""" model.train(False) model = model.to(device) out = {} for T in seq_lens: best = 1e9 for _ in range(reps): x = torch.randint(0, vocab, (1, 1), device=device) if device.type == "mps": torch.mps.synchronize() t0 = time.time() if step_api and hasattr(model, "init_state"): state = model.init_state(1, device) tok = x for _ in range(T): logits, state = model.step(tok, state) # logits [B,1,V] tok = logits[:, -1:].argmax(-1) # -> [B,1] (next token) else: seq = x for _ in range(T): logits = model(seq) tok = logits[:, -1:].argmax(-1) seq = torch.cat([seq, tok], dim=1) if device.type == "mps": torch.mps.synchronize() best = min(best, time.time() - t0) out[T] = best return out